Cleanup the code in TestRequestSender and add more tests

src/Microsoft.TestPlatform.CommunicationUtilities/TestRequestSender.cs

@@ -237,7 +237,8 @@ public void CheckVersionWithTestHost()
         // Test host sends back the lower number of the two. So the highest protocol version, that both sides support is used.
         // Error case: test host can send a protocol error if it cannot find a supported version
         var protocolNegotiated = new ManualResetEvent(false);
-        EventHandler<MessageReceivedEventArgs> onMessageReceived = (sender, args) =>
+
+EventHandler<MessageReceivedEventArgs> onMessageReceived = (sender, args) =>
         {
             var message = _dataSerializer.DeserializeMessage(args.Data!);
 
@@ -318,8 +319,8 @@ public void DiscoverTests(DiscoveryCriteria discoveryCriteria, ITestDiscoveryEve
         // When testhost disconnects, it normally means there was an error in the testhost and it exited unexpectedly.
         // But when it was us who aborted the run and killed the testhost, we don't want to wait for it to report error, because there won't be any.
         if (!TrySetupMessageReceiver(
-            onMessageReceived: (sender, args) => OnDiscoveryMessageReceived(discoveryEventsHandler, args),
-            disconnectedEventArgs => OnDiscoveryAbort(discoveryEventsHandler, disconnectedEventArgs.Error, getClientError: !_isDiscoveryAborted)))
+            onMessageReceived: (_, args) => OnDiscoveryMessageReceived(discoveryEventsHandler, args),
+            onDisconnected: disconnectedEventArgs => OnDiscoveryAbort(discoveryEventsHandler, disconnectedEventArgs.Error, getClientError: !_isDiscoveryAborted)))
         {
             return;
         }
@@ -372,8 +373,8 @@ public void StartTestRun(TestRunCriteriaWithSources runCriteria, IInternalTestRu
         _messageEventHandler = eventHandler;
 
         if (!TrySetupMessageReceiver(
-            onMessageReceived: (sender, args) => OnExecutionMessageReceived(args, eventHandler),
-            onDisconnected: (disconnectedEventArgs) => OnTestRunAbort(eventHandler, disconnectedEventArgs.Error, true)))
+            onMessageReceived: (_, args) => OnExecutionMessageReceived(args, eventHandler),
+            onDisconnected: disconnectedEventArgs => OnTestRunAbort(eventHandler, disconnectedEventArgs.Error, true)))
         {
             return;
         }
@@ -415,8 +416,8 @@ public void StartTestRun(TestRunCriteriaWithTests runCriteria, IInternalTestRunE
         _messageEventHandler = eventHandler;
 
         if (!TrySetupMessageReceiver(
-            onMessageReceived: (sender, args) => OnExecutionMessageReceived(args, eventHandler),
-            onDisconnected: (disconnectedEventArgs) => OnTestRunAbort(eventHandler, disconnectedEventArgs.Error, true)))
+            onMessageReceived: (_, args) => OnExecutionMessageReceived(args, eventHandler),
+            onDisconnected: disconnectedEventArgs => OnTestRunAbort(eventHandler, disconnectedEventArgs.Error, true)))
         {
             return;
         }

test/Microsoft.TestPlatform.CommunicationUtilities.UnitTests/TestRequestSenderTests.cs

@@ -8,6 +8,7 @@
 using System.Net;
 using System.Net.Sockets;
 using System.Threading;
+using System.Threading.Tasks;
 
 using Microsoft.VisualStudio.TestPlatform.CommunicationUtilities;
 using Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.Interfaces;
@@ -758,6 +759,52 @@ public void StartTestRunShouldNotifyErrorLogMessageIfClientDisconnectedWithClien
         _mockExecutionEventsHandler.Verify(eh => eh.HandleLogMessage(TestMessageLevel.Error, It.Is<string>(s => s.Contains(expectedErrorMessage))), Times.Once);
     }
 
+    [TestMethod]
+    public void StartTestRunShouldNotifyExecutionCompleteIfClientDisconnectedBeforeRun()
+    {
+        SetupOperationAbortedPayload();
+        SetupFakeCommunicationChannel();
+
+        RaiseClientDisconnectedEvent();
+
+        _testRequestSender.StartTestRun(_testRunCriteriaWithSources, _mockExecutionEventsHandler.Object);
+
+        _mockExecutionEventsHandler.Verify(eh => eh.HandleTestRunComplete(It.Is<TestRunCompleteEventArgs>(t => t.IsAborted), null, null, null), Times.Once);
+        _mockExecutionEventsHandler.Verify(eh => eh.HandleRawMessage("SerializedAbortedPayload"), Times.Once);
+    }
+
+    [TestMethod]
+    public void StartTestRunWithTestsShouldNotifyExecutionCompleteIfClientDisconnectedBeforeRun()
+    {
+        var runCriteria = new TestRunCriteriaWithTests(new TestCase[2], "runsettings", null, null!);
+        SetupOperationAbortedPayload();
+        SetupFakeCommunicationChannel();
+
+        RaiseClientDisconnectedEvent();
+
+        _testRequestSender.StartTestRun(runCriteria, _mockExecutionEventsHandler.Object);
+
+        _mockExecutionEventsHandler.Verify(eh => eh.HandleTestRunComplete(It.Is<TestRunCompleteEventArgs>(t => t.IsAborted), null, null, null), Times.Once);
+        _mockExecutionEventsHandler.Verify(eh => eh.HandleRawMessage("SerializedAbortedPayload"), Times.Once);
+    }
+
+    [TestMethod]
+    public async Task StartTestRunWithTestsShouldNotifyExecutionCompleteIfClientDisconnectedBeforeRunInAThreadSafeWay()
+    {
+        var runCriteria = new TestRunCriteriaWithTests(new TestCase[2], "runsettings", null, null!);
+        SetupOperationAbortedPayload();
+        SetupFakeCommunicationChannel();
+
+        // Note: Even if the calls get invoked on separate threads, the request sender should send back the complete message just once.
+        var t1 = Task.Run(RaiseClientDisconnectedEvent);
+        var t2 = Task.Run(() => _testRequestSender.StartTestRun(runCriteria, _mockExecutionEventsHandler.Object));
+
+        await Task.WhenAll(t1, t2);
+
+        _mockExecutionEventsHandler.Verify(eh => eh.HandleTestRunComplete(It.Is<TestRunCompleteEventArgs>(t => t.IsAborted), null, null, null), Times.Once);
+        _mockExecutionEventsHandler.Verify(eh => eh.HandleRawMessage("SerializedAbortedPayload"), Times.Once);
+    }
+
     [TestMethod]
     public void StartTestRunShouldNotifyExecutionCompleteIfClientDisconnected()
     {