Fix another race condition

jakubch1 · jakubch1 · commit 82bb3dec8fb1 · 2020-06-16T21:15:56.000+02:00
diff --git 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;
             });
 
