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 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>(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; } /// @@ -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.