Fix SocketCommunicationManager (#2290)

Author: MarcoRossignoli

src/Microsoft.TestPlatform.CommunicationUtilities/DataCollectionTestCaseEventSender.cs

@@ -92,7 +92,7 @@ public void SendTestCaseStart(TestCaseStartEventArgs e)
             this.communicationManager.SendMessage(MessageType.DataCollectionTestStart, e);
 
             var message = this.communicationManager.ReceiveMessage();
-            if (message.MessageType != MessageType.DataCollectionTestStartAck)
+            if (message != null && message.MessageType != MessageType.DataCollectionTestStartAck)
             {
                 if (EqtTrace.IsErrorEnabled)
                 {
@@ -108,8 +108,7 @@ public Collection<AttachmentSet> SendTestCaseEnd(TestCaseEndEventArgs e)
             this.communicationManager.SendMessage(MessageType.DataCollectionTestEnd, e);
 
             var message = this.communicationManager.ReceiveMessage();
-
-            if (message.MessageType == MessageType.DataCollectionTestEndResult)
+            if (message != null && message.MessageType == MessageType.DataCollectionTestEndResult)
             {
                 attachmentSets = this.dataSerializer.DeserializePayload<Collection<AttachmentSet>>(message);
             }

src/Microsoft.TestPlatform.CommunicationUtilities/SocketCommunicationManager.cs

@@ -285,7 +285,12 @@ public void SendRawMessage(string rawMessage)
         public Message ReceiveMessage()
         {
             var rawMessage = this.ReceiveRawMessage();
-            return this.dataSerializer.DeserializeMessage(rawMessage);
+            if (!string.IsNullOrEmpty(rawMessage))
+            {
+                return this.dataSerializer.DeserializeMessage(rawMessage);
+            }
+
+            return null;
         }
 
         /// <summary>
@@ -317,7 +322,7 @@ public string ReceiveRawMessage()
             lock (this.receiveSyncObject)
             {
                 // Reading message on binaryreader is not thread-safe
-                return this.binaryReader.ReadString();
+                return this.binaryReader?.ReadString();
             }
         }
 

test/Microsoft.TestPlatform.CommunicationUtilities.PlatformTests/SocketCommunicationManagerTests.cs

@@ -297,6 +297,22 @@ public void SocketPollShouldNotHangServerClientCommunication()
             Assert.IsTrue(true);
         }
 
+        [TestMethod]
+        public async Task ReceiveRawMessageNotConnectedSocketShouldReturnNull()
+        {
+            var peer = new SocketCommunicationManager();
+            Assert.IsNull(peer.ReceiveRawMessage());
+            Assert.IsNull(await peer.ReceiveRawMessageAsync(CancellationToken.None));
+        }
+
+        [TestMethod]
+        public async Task ReceiveMessageNotConnectedSocketShouldReturnNull()
+        {
+            var peer = new SocketCommunicationManager();
+            Assert.IsNull(peer.ReceiveMessage());
+            Assert.IsNull(await peer.ReceiveMessageAsync(CancellationToken.None));
+        }
+
         private static void SendData(ICommunicationManager communicationManager)
         {
             // Having less than the buffer size in SocketConstants.BUFFERSIZE.