Ensure dotnet test messages are received in the right timings/modes (#50727)

Author: Youssef1313

src/Cli/dotnet/Commands/CliCommandStrings.resx

@@ -2682,4 +2682,13 @@ Proceed?</value>
   <data name="UnexpectedHelpMessage" xml:space="preserve">
     <value>An unexpected help-related message was received when `--help` wasn't used.</value>
   </data>
-</root>
+  <data name="MismatchingHandshakeInfo" xml:space="preserve">
+    <value>A new handshake '{0}' was received for the test application that doesn't match the previous handshake '{1}'.</value>
+  </data>
+  <data name="UnexpectedMessageInHelpMode" xml:space="preserve">
+    <value>A message of type '{0}' was received in help mode, which is not expected.</value>
+  </data>
+  <data name="UnexpectedMessageWithoutHandshake" xml:space="preserve">
+    <value>A message of type '{0}' was received before we got any handshakes, which is unexpected.</value>
+  </data>
+</root>

src/Cli/dotnet/Commands/Test/MTP/IPC/Models/HandshakeMessage.cs

@@ -3,4 +3,4 @@
 
 namespace Microsoft.DotNet.Cli.Commands.Test.IPC.Models;
 
-internal sealed record HandshakeMessage(Dictionary<byte, string>? Properties) : IRequest, IResponse;
+internal sealed record HandshakeMessage(Dictionary<byte, string> Properties) : IRequest, IResponse;

src/Cli/dotnet/Commands/Test/MTP/IPC/Serializers/HandshakeMessageSerializer.cs

@@ -1,8 +1,6 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-#nullable disable
-
 using System.Diagnostics;
 using Microsoft.DotNet.Cli.Commands.Test.IPC.Models;
 

src/Cli/dotnet/Commands/Test/MTP/TestApplication.cs

@@ -1,11 +1,10 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-#nullable disable
-
 using System.Diagnostics;
 using System.Globalization;
 using System.IO.Pipes;
+using System.Threading;
 using Microsoft.DotNet.Cli.Commands.Test.IPC;
 using Microsoft.DotNet.Cli.Commands.Test.IPC.Models;
 using Microsoft.DotNet.Cli.Commands.Test.IPC.Serializers;
@@ -28,9 +27,7 @@ internal sealed class TestApplication(
     private readonly List<string> _outputData = [];
     private readonly List<string> _errorData = [];
     private readonly string _pipeName = NamedPipeServer.GetPipeName(Guid.NewGuid().ToString("N"));
-    private readonly CancellationTokenSource _cancellationToken = new();
 
-    private Task _testAppPipeConnectionLoop;
     private readonly List<NamedPipeServer> _testAppPipeConnections = [];
     private readonly Dictionary<NamedPipeServer, HandshakeMessage> _handshakes = new();
 
@@ -45,17 +42,26 @@ public async Task<int> RunAsync()
         // Consider throwing an exception if it's called more than once.
         var processStartInfo = CreateProcessStartInfo();
 
-        _testAppPipeConnectionLoop = Task.Run(async () => await WaitConnectionAsync(_cancellationToken.Token), _cancellationToken.Token);
-        var testProcessExitCode = await StartProcess(processStartInfo);
+        var cancellationTokenSource = new CancellationTokenSource();
+        var cancellationToken = cancellationTokenSource.Token;
+        Task? testAppPipeConnectionLoop = null;
 
-        WaitOnTestApplicationPipeConnectionLoop();
+        try
+        {
+            testAppPipeConnectionLoop = Task.Run(async () => await WaitConnectionAsync(cancellationToken), cancellationToken);
+            var testProcessExitCode = await StartProcess(processStartInfo);
+            if (_handler.HasMismatchingTestSessionEventCount())
+            {
+                throw new InvalidOperationException(CliCommandStrings.MissingTestSessionEnd);
+            }
 
-        if (_handler.HasMismatchingTestSessionEventCount())
+            return testProcessExitCode;
+        }
+        finally
         {
-            throw new InvalidOperationException(CliCommandStrings.MissingTestSessionEnd);
+            cancellationTokenSource.Cancel();
+            testAppPipeConnectionLoop?.Wait((int)TimeSpan.FromSeconds(30).TotalMilliseconds);
         }
-
-        return testProcessExitCode;
     }
 
     private ProcessStartInfo CreateProcessStartInfo()
@@ -147,12 +153,6 @@ private string GetArguments()
         return builder.ToString();
     }
 
-    private void WaitOnTestApplicationPipeConnectionLoop()
-    {
-        _cancellationToken.Cancel();
-        _testAppPipeConnectionLoop?.Wait((int)TimeSpan.FromSeconds(30).TotalMilliseconds);
-    }
-
     private async Task WaitConnectionAsync(CancellationToken token)
     {
         try
@@ -242,7 +242,7 @@ private Task<IResponse> OnRequest(NamedPipeServer server, IRequest request)
 
     private static string GetSupportedProtocolVersion(HandshakeMessage handshakeMessage)
     {
-        if (!handshakeMessage.Properties.TryGetValue(HandshakeMessagePropertyNames.SupportedProtocolVersions, out string protocolVersions) ||
+        if (!handshakeMessage.Properties.TryGetValue(HandshakeMessagePropertyNames.SupportedProtocolVersions, out string? protocolVersions) ||
             protocolVersions is null)
         {
             // It's not expected we hit this.
@@ -265,7 +265,7 @@ private static string GetSupportedProtocolVersion(HandshakeMessage handshakeMess
     }
 
     private static HandshakeMessage CreateHandshakeMessage(string version) =>
-        new(new Dictionary<byte, string>(capacity: 5)
+        new HandshakeMessage(new Dictionary<byte, string>(capacity: 5)
         {
             { HandshakeMessagePropertyNames.PID, Environment.ProcessId.ToString(CultureInfo.InvariantCulture) },
             { HandshakeMessagePropertyNames.Architecture, RuntimeInformation.ProcessArchitecture.ToString() },
@@ -278,7 +278,7 @@ private async Task<int> StartProcess(ProcessStartInfo processStartInfo)
     {
         Logger.LogTrace($"Starting test process with command '{processStartInfo.FileName}' and arguments '{processStartInfo.Arguments}'.");
 
-        using var process = Process.Start(processStartInfo);
+        using var process = Process.Start(processStartInfo)!;
         StoreOutputAndErrorData(process);
         await process.WaitForExitAsync();
 
@@ -400,7 +400,5 @@ public void Dispose()
                 Reporter.Error.WriteLine(messageBuilder.ToString());
             }
         }
-
-        WaitOnTestApplicationPipeConnectionLoop();
     }
 }

src/Cli/dotnet/Commands/Test/MTP/TestApplicationHandler.cs

@@ -46,17 +46,18 @@ internal void OnHandshakeReceived(HandshakeMessage handshakeMessage, bool gotSup
             return;
         }
 
+        var executionId = handshakeMessage.Properties[HandshakeMessagePropertyNames.ExecutionId];
+        var arch = handshakeMessage.Properties[HandshakeMessagePropertyNames.Architecture]?.ToLower();
+        var tfm = TargetFrameworkParser.GetShortTargetFramework(handshakeMessage.Properties[HandshakeMessagePropertyNames.Framework]);
+        var currentHandshakeInfo = (tfm, arch, executionId);
+
         if (!_handshakeInfo.HasValue)
         {
-            var executionId = handshakeMessage.Properties[HandshakeMessagePropertyNames.ExecutionId];
-            var arch = handshakeMessage.Properties[HandshakeMessagePropertyNames.Architecture]?.ToLower();
-            var tfm = TargetFrameworkParser.GetShortTargetFramework(handshakeMessage.Properties[HandshakeMessagePropertyNames.Framework]);
-
-            _handshakeInfo = (tfm, arch, executionId);
+            _handshakeInfo = currentHandshakeInfo;
         }
-        else
+        else if (_handshakeInfo.Value != currentHandshakeInfo)
         {
-            // TODO: Verify we get the same info.
+            throw new InvalidOperationException(string.Format(CliCommandStrings.MismatchingHandshakeInfo, currentHandshakeInfo, _handshakeInfo.Value));
         }
 
         var hostType = handshakeMessage.Properties[HandshakeMessagePropertyNames.HostType];
@@ -91,13 +92,14 @@ internal void OnDiscoveredTestsReceived(DiscoveredTestMessages discoveredTestMes
     {
         LogDiscoveredTests(discoveredTestMessages);
 
-        // TODO: If _handshakeInfo is null, we should error.
-        // We shouldn't be getting any discovered test messages without a previous handshake.
-
         if (_options.IsHelp)
         {
-            // TODO: Better to throw exception?
-            return;
+            throw new InvalidOperationException(string.Format(CliCommandStrings.UnexpectedMessageInHelpMode, nameof(DiscoveredTestMessages)));
+        }
+
+        if (!_handshakeInfo.HasValue)
+        {
+            throw new InvalidOperationException(string.Format(CliCommandStrings.UnexpectedMessageWithoutHandshake, nameof(DiscoveredTestMessages)));
         }
 
         foreach (var test in discoveredTestMessages.DiscoveredMessages)
@@ -112,13 +114,14 @@ internal void OnTestResultsReceived(TestResultMessages testResultMessage)
     {
         LogTestResults(testResultMessage);
 
-        // TODO: If _handshakeInfo is null, we should error.
-        // We shouldn't be getting any test result messages without a previous handshake.
-
         if (_options.IsHelp)
         {
-            // TODO: Better to throw exception?
-            return;
+            throw new InvalidOperationException(string.Format(CliCommandStrings.UnexpectedMessageInHelpMode, nameof(TestResultMessages)));
+        }
+
+        if (!_handshakeInfo.HasValue)
+        {
+            throw new InvalidOperationException(string.Format(CliCommandStrings.UnexpectedMessageWithoutHandshake, nameof(TestResultMessages)));
         }
 
         var handshakeInfo = _handshakeInfo.Value;
@@ -158,13 +161,14 @@ internal void OnFileArtifactsReceived(FileArtifactMessages fileArtifactMessages)
     {
         LogFileArtifacts(fileArtifactMessages);
 
-        // TODO: If _handshakeInfo is null, we should error.
-        // We shouldn't be getting any file artifact messages without a previous handshake.
-
         if (_options.IsHelp)
         {
-            // TODO: Better to throw exception?
-            return;
+            throw new InvalidOperationException(string.Format(CliCommandStrings.UnexpectedMessageInHelpMode, nameof(FileArtifactMessages)));
+        }
+
+        if (!_handshakeInfo.HasValue)
+        {
+            throw new InvalidOperationException(string.Format(CliCommandStrings.UnexpectedMessageWithoutHandshake, nameof(FileArtifactMessages)));
         }
 
         var handshakeInfo = _handshakeInfo.Value;
@@ -183,8 +187,12 @@ internal void OnSessionEventReceived(TestSessionEvent sessionEvent)
         {
             LogSessionEvent(sessionEvent);
 
-            // TODO: If _handshakeInfo is null, we should error.
-            // We shouldn't be getting any session event messages without a previous handshake.
+            // TODO: Validate if we should get this message in help mode or not.
+
+            if (!_handshakeInfo.HasValue)
+            {
+                throw new InvalidOperationException(string.Format(CliCommandStrings.UnexpectedMessageWithoutHandshake, nameof(DiscoveredTestMessages)));
+            }
 
             if (sessionEvent.SessionType == SessionEventTypes.TestSessionStart)
             {