From 11908955ca7a09035dd3a6ddc7ab3e2273edc904 Mon Sep 17 00:00:00 2001
From: Jakub Chocholowicz <jakub.chocholowicz@gmail.com>
Date: Tue, 16 Jun 2020 19:23:06 +0200
Subject: [PATCH] Fix race condition in test

---
 .../DesignMode/DesignModeClientTests.cs                   | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/test/Microsoft.TestPlatform.Client.UnitTests/DesignMode/DesignModeClientTests.cs b/test/Microsoft.TestPlatform.Client.UnitTests/DesignMode/DesignModeClientTests.cs
index e34ec2335c..4c8c962e92 100644
--- a/test/Microsoft.TestPlatform.Client.UnitTests/DesignMode/DesignModeClientTests.cs
+++ b/test/Microsoft.TestPlatform.Client.UnitTests/DesignMode/DesignModeClientTests.cs
@@ -460,8 +460,16 @@ public void DesignModeClientConnectShouldCallRequestManagerForFinalizationStart(
             this.mockCommunicationManager.Setup(cm => cm.WaitForServerConnection(It.IsAny<int>())).Returns(true);
             this.mockCommunicationManager.SetupSequence(cm => cm.ReceiveMessage()).Returns(startFinalization);
 
+            this.mockTestRequestManager
+                .Setup(
+                    rm => rm.FinalizeMultiTestRun(
+                        It.IsAny<MultiTestRunFinalizationPayload>(),
+                        It.IsAny<IMultiTestRunFinalizationEventsHandler>()))
+                .Callback(() => complateEvent.Set());
+
             this.designModeClient.ConnectToClientAndProcessRequests(PortNumber, this.mockTestRequestManager.Object);
 
+            Assert.IsTrue(this.complateEvent.WaitOne(Timeout), "Finalization not completed.");
             this.mockCommunicationManager.Verify(cm => cm.SendMessage(MessageType.TestMessage, It.IsAny<TestMessagePayload>()), Times.Never);
             this.mockCommunicationManager.Verify(cm => cm.SendMessage(MessageType.MultiTestRunFinalizationComplete, It.IsAny<MultiTestRunFinalizationCompletePayload>()), Times.Never);
             this.mockTestRequestManager.Verify(rm => rm.FinalizeMultiTestRun(It.IsAny<MultiTestRunFinalizationPayload>(), It.IsAny<MultiTestRunFinalizationEventsHandler>()));