From e8e994aa009b5c5794e20edd228c41444e7eeb88 Mon Sep 17 00:00:00 2001
From: Marco Rossignoli <marco.rossignoli@gmail.com>
Date: Mon, 13 Jan 2020 13:28:59 +0100
Subject: [PATCH] Fix SocketCommunicationManager (#2290)

---
 .../DataCollectionTestCaseEventSender.cs         |  5 ++---
 .../SocketCommunicationManager.cs                |  9 +++++++--
 .../SocketCommunicationManagerTests.cs           | 16 ++++++++++++++++
 3 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/src/Microsoft.TestPlatform.CommunicationUtilities/DataCollectionTestCaseEventSender.cs b/src/Microsoft.TestPlatform.CommunicationUtilities/DataCollectionTestCaseEventSender.cs
index ead6e6fedc..0f2aba6cd3 100644
--- a/src/Microsoft.TestPlatform.CommunicationUtilities/DataCollectionTestCaseEventSender.cs
+++ b/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);
             }
diff --git a/src/Microsoft.TestPlatform.CommunicationUtilities/SocketCommunicationManager.cs b/src/Microsoft.TestPlatform.CommunicationUtilities/SocketCommunicationManager.cs
index 50d19c7687..411dcaf88b 100644
--- a/src/Microsoft.TestPlatform.CommunicationUtilities/SocketCommunicationManager.cs
+++ b/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();
             }
         }
 
diff --git a/test/Microsoft.TestPlatform.CommunicationUtilities.PlatformTests/SocketCommunicationManagerTests.cs b/test/Microsoft.TestPlatform.CommunicationUtilities.PlatformTests/SocketCommunicationManagerTests.cs
index b070989b1b..9faed62488 100644
--- a/test/Microsoft.TestPlatform.CommunicationUtilities.PlatformTests/SocketCommunicationManagerTests.cs
+++ b/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.