From 303cee86825aa1dd70b17ac62aa8062fdcdf5ba1 Mon Sep 17 00:00:00 2001
From: Jakub Chocholowicz <jakub.chocholowicz@gmail.com>
Date: Mon, 22 Jun 2020 23:19:06 +0200
Subject: [PATCH] Fix test

---
 .../MultiTestRunFinalizationManagerTests.cs         | 13 +++++++++++++
 1 file changed, 13 insertions(+)

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<object>();
+
+            int counter = 0;
+            mockEventsHandler.Setup(h => h.HandleMultiTestRunFinalizationProgress(It.IsAny<MultiTestRunFinalizationProgressEventArgs>())).Callback(() => 
+            { 
+                counter++;
+                if(counter == 6)
+                {
+                    innerTaskCompletionSource.TrySetResult(null);
+                }
+            });
+
             mockAttachmentHandler1.Setup(h => h.HandleDataCollectionAttachmentSets(It.IsAny<ICollection<AttachmentSet>>(), It.IsAny<IProgress<int>>(), It.IsAny<IMessageLogger>(), It.IsAny<CancellationToken>())).Returns((ICollection<AttachmentSet> i1, IProgress<int> 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<MultiTestRunFinalizationProgressEventArgs>()), Times.Exactly(6));
             mockEventsHandler.Verify(h => h.HandleMultiTestRunFinalizationProgress(It.Is<MultiTestRunFinalizationProgressEventArgs>(a => VerifyProgressArgsForTwoHandlers(a, 1, 25, uri1))));