Fix up test for WinPS

Author: jborean93

src/PowerShellEditorServices/Services/DebugAdapter/Handlers/LaunchAndAttachHandler.cs

@@ -427,7 +427,7 @@ await _executionService.ExecutePSCommandAsync(
             }
             catch (Exception e)
             {
-                string msg = $"Could not attach to process with ID: {processId} - {e.Message}\n{e.StackTrace}";
+                string msg = $"Could not attach to process with ID: {processId}";
                 _logger.LogError(e, msg);
                 throw new RpcErrorException(0, null, msg);
             }

test/PowerShellEditorServices.Test.E2E/DebugAdapterProtocolMessageTests.cs

@@ -532,7 +532,7 @@ public async Task CanAttachScriptWithPathMappings()
 
             string[] logStatements = ["$PSCommandPath", "after breakpoint"];
 
-            await RunWithAttachableProcess(logStatements, async (filePath, processId) =>
+            await RunWithAttachableProcess(logStatements, async (filePath, processId, runspaceId) =>
             {
                 string localParent = Path.Combine(Path.GetTempPath(), Path.GetRandomFileName());
                 string localScriptPath = Path.Combine(localParent, Path.GetFileName(filePath));
@@ -541,11 +541,11 @@ await RunWithAttachableProcess(logStatements, async (filePath, processId) =>
 
                 Task<StoppedEvent> nextStoppedTask = nextStopped;
 
-                _ = await client.Attach(
+                AttachResponse attachResponse = await client.Attach(
                     new PsesAttachRequestArguments
                     {
                         ProcessId = processId,
-                        RunspaceId = 1,
+                        RunspaceId = runspaceId,
                         PathMappings = [
                             new()
                             {
@@ -554,6 +554,7 @@ await RunWithAttachableProcess(logStatements, async (filePath, processId) =>
                             }
                         ]
                     }) ?? throw new Exception("Attach response was null.");
+                Assert.NotNull(attachResponse);
 
                 SetBreakpointsResponse setBreakpointsResponse = await client.SetBreakpoints(new SetBreakpointsArguments
                 {
@@ -585,12 +586,7 @@ await RunWithAttachableProcess(logStatements, async (filePath, processId) =>
 
                 // Wait until we hit the breakpoint
                 stoppedEvent = await nextStoppedTask;
-
-                // WinPS has a bug on attach where it doesn't send the breakpoints on the stop event.
-                string expectedReason = PsesStdioLanguageServerProcessHost.IsWindowsPowerShell
-                    ? "step"
-                    : "breakpoint";
-                Assert.Equal(expectedReason, stoppedEvent.Reason);
+                Assert.Equal("breakpoint", stoppedEvent.Reason);
 
                 // The code before the breakpoint should have already run
                 // It will contain the actual script being run
@@ -612,33 +608,26 @@ await RunWithAttachableProcess(logStatements, async (filePath, processId) =>
             });
         }
 
-        private async Task RunWithAttachableProcess(string[] logStatements, Func<string, int, Task> action)
+        private async Task RunWithAttachableProcess(string[] logStatements, Func<string, int, 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
-            // subscribed to by Debug-Runspace as a marker that it is ready to
-            // continue.
-            // Use https://github.com/PowerShell/PowerShell/pull/25788 once it
-            // is available.
+            /*
+                There is no public API in pwsh to wait for an attach event. We
+                use reflection to wait until the AvailabilityChanged event is
+                subscribed to by Debug-Runspace as a marker that it is ready to
+                continue.
+
+                We also run the test script in another runspace as WinPS'
+                Debug-Runspace will break on the first statement after the
+                attach and we want that to be the Wait-Debugger call.
+
+                We can use https://github.com/PowerShell/PowerShell/pull/25788
+                once that is merged and we are running against that version but
+                WinPS will always need this.
+            */
             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
+                param([string]$TestScript)
 
                 $debugRunspaceCmd = Get-Command Debug-Runspace -Module Microsoft.PowerShell.Utility
-                $runspace = [Runspace]::DefaultRunspace
                 $runspaceBase = [PSObject].Assembly.GetType(
                     'System.Management.Automation.Runspaces.RunspaceBase')
                 $availabilityChangedField = $runspaceBase.GetField(
@@ -648,6 +637,18 @@ private async Task RunWithAttachableProcess(string[] logStatements, Func<string,
                     throw 'Failed to get AvailabilityChanged event field'
                 }
 
+                $ps = [PowerShell]::Create()
+                $runspace = $ps.Runspace
+
+                # Wait-Debugger is needed in WinPS to sync breakpoints before
+                # running the script.
+                $null = $ps.AddCommand('Wait-Debugger').AddStatement()
+                $null = $ps.AddCommand($TestScript)
+
+                # Let the runner know what Runspace to attach to and that it
+                # is ready to run.
+                'RID: {0}' -f $runspace.Id
+
                 $start = Get-Date
                 while ($true) {
                     $subscribed = $availabilityChangedField.GetValue($runspace) |
@@ -661,10 +662,10 @@ private async Task RunWithAttachableProcess(string[] logStatements, Func<string,
                     }
                 }
 
-                # Needed to sync breakpoints on WinPS 5.1
-                Wait-Debugger
-
-                & $TestScript
+                $ps.Invoke()
+                foreach ($e in $ps.Streams.Error) {
+                    Write-Error -ErrorRecord $e
+                }
 
                 # Keep running until the runner has deleted the test script to
                 # ensure the process doesn't finish before the test does in
@@ -674,27 +675,8 @@ private async Task RunWithAttachableProcess(string[] logStatements, Func<string,
                 }
             ";
 
-            // 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 filePath = NewTestFile(GenerateLoggingScript(logStatements));
+            string encArgs = CreatePwshEncodedArgs(filePath);
             string encCommand = Convert.ToBase64String(Encoding.Unicode.GetBytes(scriptEntrypoint));
 
             ProcessStartInfo psi = new ProcessStartInfo
@@ -708,15 +690,24 @@ private async Task RunWithAttachableProcess(string[] logStatements, Func<string,
             };
             psi.EnvironmentVariables["TERM"] = "dumb";  // Avoids color/VT sequences in test output.
 
+            TaskCompletionSource<int> ridOutput = new();
+
             // Task shouldn't take longer than 30 seconds to complete.
             using CancellationTokenSource debugTaskCts = new CancellationTokenSource(TimeSpan.FromSeconds(30));
+            using CancellationTokenRegistration _ = debugTaskCts.Token.Register(ridOutput.SetCanceled);
             using Process psProc = Process.Start(psi);
             try
             {
                 psProc.OutputDataReceived += (sender, args) =>
                 {
                     if (!string.IsNullOrEmpty(args.Data))
                     {
+                        if (args.Data.StartsWith("RID: "))
+                        {
+                            int rid = int.Parse(args.Data.Substring(5));
+                            ridOutput.SetResult(rid);
+                        }
+
                         output.WriteLine("STDOUT: {0}", args.Data);
                     }
                 };
@@ -732,24 +723,18 @@ private async Task RunWithAttachableProcess(string[] logStatements, Func<string,
                 psProc.BeginErrorReadLine();
 
                 Task procExited = psProc.WaitForExitAsync(debugTaskCts.Token);
-                Task waitStart = Task.Run(async () =>
-                {
-                    while (File.Exists(startMarker))
-                    {
-                        await Task.Delay(100, debugTaskCts.Token);
-                    }
-                });
+                Task<int> waitRid = ridOutput.Task;
 
                 // Wait for the process to fail or the script to start.
-                Task finishedTask = await Task.WhenAny(waitStart, procExited);
+                Task finishedTask = await Task.WhenAny(waitRid, procExited);
                 if (finishedTask == procExited)
                 {
                     await procExited;
                     Assert.Fail("The attached process exited before the PowerShell entrypoint could start.");
                 }
-                await waitStart;
+                int rid = await waitRid;
 
-                Task debugTask = action(filePath, psProc.Id);
+                Task debugTask = action(filePath, psProc.Id, rid);
                 finishedTask = await Task.WhenAny(procExited, debugTask);
                 if (finishedTask == procExited)
                 {
@@ -773,5 +758,30 @@ private async Task RunWithAttachableProcess(string[] logStatements, Func<string,
                 throw;
             }
         }
+
+        private static string CreatePwshEncodedArgs(params string[] args)
+        {
+            // Only way to pass args to -EncodedCommand is to use CLIXML with
+            // -EncodedArguments. Not pretty but the structure isn't too
+            // complex and saves us trying to embed/escape strings in a script.
+            string clixmlNamespace = "http://schemas.microsoft.com/powershell/2004/04";
+            string clixml = 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),
+                            args.Select(s => new XElement(XName.Get("S", clixmlNamespace), s))
+                        )
+                ))).ToString(SaveOptions.DisableFormatting);
+
+            return Convert.ToBase64String(Encoding.Unicode.GetBytes(clixml));
+        }
     }
 }