Fix HeartbeatTests flakiness (#2268)

dotnet · Jan 23, 2018 · ee1c5c4 · ee1c5c4
1 parent 526dfdb
commit ee1c5c4
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/test/Kestrel.Core.Tests/HeartbeatTests.cs b/test/Kestrel.Core.Tests/HeartbeatTests.cs
@@ -30,19 +30,21 @@ public void BlockedHeartbeatDoesntCauseOverlapsAndIsLoggedAsError()
             var kestrelTrace = new Mock<IKestrelTrace>();
             var handlerMre = new ManualResetEventSlim();
             var traceMre = new ManualResetEventSlim();
+            var onHeartbeatTasks = new Task[2];
 
             heartbeatHandler.Setup(h => h.OnHeartbeat(systemClock.UtcNow)).Callback(() => handlerMre.Wait());
             debugger.Setup(d => d.IsAttached).Returns(false);
             kestrelTrace.Setup(t => t.HeartbeatSlow(Heartbeat.Interval, systemClock.UtcNow)).Callback(() => traceMre.Set());
 
             using (var heartbeat = new Heartbeat(new[] { heartbeatHandler.Object }, systemClock, debugger.Object, kestrelTrace.Object))
             {
-                Task.Run(() => heartbeat.OnHeartbeat());
-                Task.Run(() => heartbeat.OnHeartbeat());
+                onHeartbeatTasks[0] = Task.Run(() => heartbeat.OnHeartbeat());
+                onHeartbeatTasks[1] = Task.Run(() => heartbeat.OnHeartbeat());
                 Assert.True(traceMre.Wait(TimeSpan.FromSeconds(10)));
             }
 
             handlerMre.Set();
+            Task.WaitAll(onHeartbeatTasks);
 
             heartbeatHandler.Verify(h => h.OnHeartbeat(systemClock.UtcNow), Times.Once());
             kestrelTrace.Verify(t => t.HeartbeatSlow(Heartbeat.Interval, systemClock.UtcNow), Times.Once());
@@ -57,19 +59,21 @@ public void BlockedHeartbeatIsNotLoggedAsErrorIfDebuggerAttached()
             var kestrelTrace = new Mock<IKestrelTrace>();
             var handlerMre = new ManualResetEventSlim();
             var traceMre = new ManualResetEventSlim();
+            var onHeartbeatTasks = new Task[2];
 
             heartbeatHandler.Setup(h => h.OnHeartbeat(systemClock.UtcNow)).Callback(() => handlerMre.Wait());
             debugger.Setup(d => d.IsAttached).Returns(true);
             kestrelTrace.Setup(t => t.HeartbeatSlow(Heartbeat.Interval, systemClock.UtcNow)).Callback(() => traceMre.Set());
 
             using (var heartbeat = new Heartbeat(new[] { heartbeatHandler.Object }, systemClock, debugger.Object, kestrelTrace.Object))
             {
-                Task.Run(() => heartbeat.OnHeartbeat());
-                Task.Run(() => heartbeat.OnHeartbeat());
+                onHeartbeatTasks[0] = Task.Run(() => heartbeat.OnHeartbeat());
+                onHeartbeatTasks[1] = Task.Run(() => heartbeat.OnHeartbeat());
                 Assert.False(traceMre.Wait(TimeSpan.FromSeconds(10)));
             }
 
             handlerMre.Set();
+            Task.WaitAll(onHeartbeatTasks);
 
             heartbeatHandler.Verify(h => h.OnHeartbeat(systemClock.UtcNow), Times.Once());
             kestrelTrace.Verify(t => t.HeartbeatSlow(Heartbeat.Interval, systemClock.UtcNow), Times.Never());