From 82bb3dec8fb1eff8cc0c9c06557674eaf5b5c4a1 Mon Sep 17 00:00:00 2001
From: Jakub Chocholowicz <jakub.chocholowicz@gmail.com>
Date: Tue, 16 Jun 2020 21:15:56 +0200
Subject: [PATCH] Fix another race condition

---
 .../MultiTestRunFinalizationManagerTests.cs   | 27 ++++++++++++-------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/test/Microsoft.TestPlatform.CrossPlatEngine.UnitTests/MultiTestRunFinalization/MultiTestRunFinalizationManagerTests.cs b/test/Microsoft.TestPlatform.CrossPlatEngine.UnitTests/MultiTestRunFinalization/MultiTestRunFinalizationManagerTests.cs
index 5a681f38c4..094d26fa05 100644
--- a/test/Microsoft.TestPlatform.CrossPlatEngine.UnitTests/MultiTestRunFinalization/MultiTestRunFinalizationManagerTests.cs
+++ b/test/Microsoft.TestPlatform.CrossPlatEngine.UnitTests/MultiTestRunFinalization/MultiTestRunFinalizationManagerTests.cs
@@ -392,20 +392,27 @@ public async Task FinalizeMultiTestRunAsync_ShouldReturnNullThroughEventsHandler
 
             mockAttachmentHandler1.Setup(h => h.HandleDataCollectionAttachmentSets(It.IsAny<ICollection<AttachmentSet>>(), It.IsAny<CancellationToken>())).Returns((ICollection<AttachmentSet> i1, CancellationToken cancellation) =>
             {
-                for (int i = 0; i < 1000; ++i)
-                {                    
-                    Task.Delay(100);
-                    Console.WriteLine($"Iteration: {i}");
+                try
+                {
+                    for (int i = 0; i < 1000; ++i)
+                    {
+                        Task.Delay(100).Wait();
+                        Console.WriteLine($"Iteration: {i}");
 
-                    if (cancellation.IsCancellationRequested) break;
+                        cancellation.ThrowIfCancellationRequested();
 
-                    if (i == 3)
-                    {
-                        cancellationTokenSource.Cancel();
+                        if (i == 3)
+                        {
+                            cancellationTokenSource.Cancel();
+                            Task.Delay(500).Wait();
+                        }
                     }
                 }
-
-                innerTaskCompletionSource.TrySetResult(null);
+                finally
+                {
+                    innerTaskCompletionSource.TrySetResult(null);
+                }
+                
                 return outputAttachments;
             });