Fix another race condition

microsoft · Jun 16, 2020 · 82bb3de · 82bb3de
1 parent 1190895
commit 82bb3de
Showing 1 changed file with 17 additions and 10 deletions.
diff --git a/...rossPlatEngine.UnitTests/MultiTestRunFinalization/MultiTestRunFinalizationManagerTests.cs b/...rossPlatEngine.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;
             });