diff --git a/test/Microsoft.TestPlatform.CrossPlatEngine.UnitTests/MultiTestRunFinalization/MultiTestRunFinalizationManagerTests.cs b/test/Microsoft.TestPlatform.CrossPlatEngine.UnitTests/MultiTestRunFinalization/MultiTestRunFinalizationManagerTests.cs index 70ceea583b..0b5c678a8a 100644 --- a/test/Microsoft.TestPlatform.CrossPlatEngine.UnitTests/MultiTestRunFinalization/MultiTestRunFinalizationManagerTests.cs +++ b/test/Microsoft.TestPlatform.CrossPlatEngine.UnitTests/MultiTestRunFinalization/MultiTestRunFinalizationManagerTests.cs @@ -533,6 +533,18 @@ public async Task FinalizeMultiTestRunAsync_ShouldReturnProperlySendProgressEven new AttachmentSet(new Uri(uri2), "uri2_output") }; + var innerTaskCompletionSource = new TaskCompletionSource(); + + int counter = 0; + mockEventsHandler.Setup(h => h.HandleMultiTestRunFinalizationProgress(It.IsAny())).Callback(() => + { + counter++; + if(counter == 6) + { + innerTaskCompletionSource.TrySetResult(null); + } + }); + mockAttachmentHandler1.Setup(h => h.HandleDataCollectionAttachmentSets(It.IsAny>(), It.IsAny>(), It.IsAny(), It.IsAny())).Returns((ICollection i1, IProgress progress, IMessageLogger logger, CancellationToken cancellation) => { progress.Report(25); @@ -555,6 +567,7 @@ public async Task FinalizeMultiTestRunAsync_ShouldReturnProperlySendProgressEven await manager.FinalizeMultiTestRunAsync(mockRequestData.Object, inputAttachments, mockEventsHandler.Object, CancellationToken.None); // assert + await innerTaskCompletionSource.Task; VerifyCompleteEvent(false, false, outputAttachments1[0], outputAttachments2[0]); mockEventsHandler.Verify(h => h.HandleMultiTestRunFinalizationProgress(It.IsAny()), Times.Exactly(6)); mockEventsHandler.Verify(h => h.HandleMultiTestRunFinalizationProgress(It.Is(a => VerifyProgressArgsForTwoHandlers(a, 1, 25, uri1))));