From deb9770f9e638c936fb5613a4d6a636af5dcdb98 Mon Sep 17 00:00:00 2001
From: Joey Robichaud <jorobich@microsoft.com>
Date: Fri, 24 Jul 2020 09:25:14 -0700
Subject: [PATCH] Fix getting ProcessStartInfo for tests

---
 src/OmniSharp.DotNetTest/VSTestManager.cs | 54 +++++++++++++++++------
 1 file changed, 41 insertions(+), 13 deletions(-)

diff --git a/src/OmniSharp.DotNetTest/VSTestManager.cs b/src/OmniSharp.DotNetTest/VSTestManager.cs
index 125943a0db..d82d736b00 100644
--- a/src/OmniSharp.DotNetTest/VSTestManager.cs
+++ b/src/OmniSharp.DotNetTest/VSTestManager.cs
@@ -157,8 +157,7 @@ public override async Task<GetTestStartInfoResponse> GetTestStartInfoAsync(strin
                     RunSettings = LoadRunSettingsOrDefault(runSettings, targetFrameworkVersion)
                 });
 
-            var message = ReadMessage();
-            var testStartInfo = message.DeserializePayload<TestProcessStartInfo>();
+            var testStartInfo = await GetTestRunnerProcessStartInfo(testCases, debuggingEnabled: false, runSettings, targetFrameworkVersion, cancellationToken);
 
             return new GetTestStartInfoResponse
             {
@@ -176,26 +175,55 @@ public override async Task<DebugTestGetStartInfoResponse> DebugGetStartInfoAsync
             VerifyTestFramework(testFrameworkName);
 
             var testCases = await DiscoverTestsAsync(methodNames, runSettings, targetFrameworkVersion, cancellationToken);
+            var testStartInfo = await GetTestRunnerProcessStartInfo(testCases, debuggingEnabled: true, runSettings, targetFrameworkVersion, cancellationToken);
 
+            return new DebugTestGetStartInfoResponse
+            {
+                FileName = testStartInfo.FileName,
+                Arguments = testStartInfo.Arguments,
+                WorkingDirectory = testStartInfo.WorkingDirectory,
+                EnvironmentVariables = testStartInfo.EnvironmentVariables,
+                Succeeded = true
+            };
+        }
+
+        private async Task<TestProcessStartInfo> GetTestRunnerProcessStartInfo(TestCase[] testCases, bool debuggingEnabled, string runSettings, string targetFrameworkVersion, CancellationToken cancellationToken)
+        {
             SendMessage(MessageType.GetTestRunnerProcessStartInfoForRunSelected,
                 new
                 {
                     TestCases = testCases,
-                    DebuggingEnabled = true,
                     RunSettings = LoadRunSettingsOrDefault(runSettings, targetFrameworkVersion)
                 });
 
-            var message = await ReadMessageAsync(cancellationToken);
-            var startInfo = message.DeserializePayload<TestProcessStartInfo>();
-
-            return new DebugTestGetStartInfoResponse
+            bool done = false;
+            TestProcessStartInfo startInfo = null;
+            while (!done)
             {
-                FileName = startInfo.FileName,
-                Arguments = startInfo.Arguments,
-                WorkingDirectory = startInfo.WorkingDirectory,
-                EnvironmentVariables = startInfo.EnvironmentVariables,
-                Succeeded = true
-            };
+                var (success, message) = await TryReadMessageAsync(cancellationToken);
+                if (!success)
+                {
+                    return startInfo;
+                }
+
+                switch (message.MessageType)
+                {
+                    case MessageType.TestMessage:
+                        EmitTestMessage(message.DeserializePayload<TestMessagePayload>());
+                        break;
+
+                    case MessageType.CustomTestHostLaunch:
+                        startInfo = message.DeserializePayload<TestProcessStartInfo>();
+                        done = true;
+                        break;
+
+                    case MessageType.ExecutionComplete:
+                        done = true;
+                        break;
+                }
+            }
+
+            return startInfo;
         }
 
         public override async Task DebugLaunchAsync(CancellationToken cancellationToken)