Add wait to test before attaching

Author: jborean93

test/PowerShellEditorServices.Test.E2E/DebugAdapterProtocolMessageTests.cs

@@ -9,6 +9,7 @@
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Xml.Linq;
 using Microsoft.PowerShell.EditorServices.Handlers;
 using Nerdbank.Streams;
 using OmniSharp.Extensions.DebugAdapter.Client;
@@ -608,6 +609,13 @@ await RunWithAttachableProcess(logStatements, async (filePath, processId) =>
 
         private async Task RunWithAttachableProcess(string[] logStatements, Func<string, int, Task> action)
         {
+            // Ensures that we don't try and attach until the PowerShell proc
+            // has fully started and is ready to accept the attach request.
+            string startMarker = Path.Combine(Path.GetTempPath(), Path.GetRandomFileName());
+            File.WriteAllText(startMarker, "");
+
+            string filePath = NewTestFile(GenerateLoggingScript(logStatements));
+
             // PowerShell has no public API for waiting when Debug-Runspace
             // has attached itself to a Runspace. We use this reflection
             // hackery to wait until the AvailabilityChanged event is
@@ -616,8 +624,14 @@ private async Task RunWithAttachableProcess(string[] logStatements, Func<string,
             // Use https://github.com/PowerShell/PowerShell/pull/25788 once it
             // is available.
             string scriptEntrypoint = @"
+                param([string]$StartMarker, [string]$TestScript)
+
                 $ErrorActionPreference = 'Stop'
 
+                # Removing this file tells the runner that the script has
+                # started and pwsh is ready to accept the attach request.
+                Remove-Item -LiteralPath $StartMarker -Force
+
                 $debugRunspaceCmd = Get-Command Debug-Runspace -Module Microsoft.PowerShell.Utility
                 $runspace = [Runspace]::DefaultRunspace
                 $runspaceBase = [PSObject].Assembly.GetType(
@@ -644,22 +658,44 @@ private async Task RunWithAttachableProcess(string[] logStatements, Func<string,
 
                 # Needed to sync breakpoints on WinPS 5.1
                 Wait-Debugger
-            ";
 
-            string filePath = NewTestFile(GenerateLoggingScript(logStatements));
-            scriptEntrypoint += Environment.NewLine + $"& '{filePath}'";
+                & $TestScript
 
-            // This allows us to ensure the process won't end after the script
-            // has run. Without this we run the risk that the process ends
-            // before the debug task completes causing a test failure.
-            scriptEntrypoint += Environment.NewLine + $"while (Test-Path -Path '{filePath}') {{ Start-Sleep -Seconds 1 }}";
+                # Keep running until the runner has deleted the test script to
+                # ensure the process doesn't finish before the test does in
+                # normal circumstances.
+                while (Test-Path -Path $TestScript) {
+                    Start-Sleep -Seconds 1
+                }
+            ";
 
+            // Only way to pass args to -EncodedCommand is to use CLIXML with
+            // -EncodedArguments. Not pretty but the structure isn't too
+            // complex.
+            string clixmlNamespace = "http://schemas.microsoft.com/powershell/2004/04";
+            string clixmlArg = new XDocument(
+                new XDeclaration("1.0", "utf-16", "yes"),
+                new XElement(XName.Get("Objs", clixmlNamespace),
+                    new XAttribute("Version", "1.1.0.1"),
+                    new XElement(XName.Get("Obj", clixmlNamespace),
+                        new XAttribute("RefId", "0"),
+                        new XElement(XName.Get("TN", clixmlNamespace),
+                            new XAttribute("RefId", "0"),
+                            new XElement(XName.Get("T", clixmlNamespace), "System.Collections.ArrayList"),
+                            new XElement(XName.Get("T", clixmlNamespace), "System.Object")
+                        ),
+                        new XElement(XName.Get("LST", clixmlNamespace),
+                            new XElement(XName.Get("S", clixmlNamespace), startMarker),
+                            new XElement(XName.Get("S", clixmlNamespace), filePath)
+                        )
+                ))).ToString(SaveOptions.DisableFormatting);
+            string encArgs = Convert.ToBase64String(Encoding.Unicode.GetBytes(clixmlArg));
             string encCommand = Convert.ToBase64String(Encoding.Unicode.GetBytes(scriptEntrypoint));
 
             ProcessStartInfo psi = new ProcessStartInfo
             {
                 FileName = PsesStdioLanguageServerProcessHost.PwshExe,
-                Arguments = $"-NoLogo -NoProfile -EncodedCommand {encCommand}",
+                Arguments = $"-NoLogo -NoProfile -EncodedCommand {encCommand} -EncodedArguments {encArgs}",
                 RedirectStandardOutput = true,
                 RedirectStandardError = true,
                 UseShellExecute = false,
@@ -691,13 +727,29 @@ private async Task RunWithAttachableProcess(string[] logStatements, Func<string,
                 psProc.BeginErrorReadLine();
 
                 Task procExited = psProc.WaitForExitAsync(debugTaskCts.Token);
-                Task debugTask = action(filePath, psProc.Id);
+                Task waitStart = Task.Run(async () =>
+                {
+                    while (File.Exists(startMarker))
+                    {
+                        await Task.Delay(100, debugTaskCts.Token);
+                    }
+                });
+
+                // Wait for the process to fail or the script to start.
+                Task finishedTask = await Task.WhenAny(waitStart, procExited);
+                if (finishedTask == procExited)
+                {
+                    await procExited;
+                    Assert.Fail("The attached process exited before the PowerShell entrypoint could start.");
+                }
+                await waitStart;
 
-                Task finishedTask = await Task.WhenAny(procExited, debugTask);
+                Task debugTask = action(filePath, psProc.Id);
+                finishedTask = await Task.WhenAny(procExited, debugTask);
                 if (finishedTask == procExited)
                 {
                     await procExited;
-                    Assert.Fail("Attached process exited before the debug task completed.");
+                    Assert.Fail("Attached process exited before the script could start.");
                 }
 
                 await debugTask;