Fix SocketCommunicationManager (#2290)

microsoft · Jan 13, 2020 · e8e994a · e8e994a
1 parent 3ae0859
commit e8e994a
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 5 deletions.
diff --git 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
@@ -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/...soft.TestPlatform.CommunicationUtilities.PlatformTests/SocketCommunicationManagerTests.cs b/...soft.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.