Improve E2E Test Fixtures to be less flaky

src/PowerShellEditorServices/Server/PsesDebugServer.cs

@@ -120,7 +120,8 @@ public async Task StartAsync()
                         response.SupportsDelayedStackTraceLoading = true;
 
                         return Task.CompletedTask;
-                    });
+                    })
+                ;
             }).ConfigureAwait(false);
         }
 

test/PowerShellEditorServices.Test.E2E/Hosts/DebugOutputStream.cs

@@ -0,0 +1,44 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+#nullable enable
+
+using System.Diagnostics;
+using System.IO;
+using System.Text;
+using Nerdbank.Streams;
+
+namespace PowerShellEditorServices.Test.E2E;
+
+/// <summary>
+/// A stream that logs all data read and written to the debug stream which is visible in the debug console when a
+/// debugger is attached.
+/// </summary>
+internal class DebugOutputStream : MonitoringStream
+{
+    public DebugOutputStream(Stream? underlyingStream)
+    : base(underlyingStream ?? new MemoryStream())
+    {
+        DidRead += (_, segment) =>
+        {
+            if (segment.Array is null) { return; }
+            LogData("⬅️", segment.Array, segment.Offset, segment.Count);
+        };
+
+        DidWrite += (_, segment) =>
+        {
+            if (segment.Array is null) { return; }
+            LogData("➡️", segment.Array, segment.Offset, segment.Count);
+        };
+
+    }
+
+    private static void LogData(string header, byte[] buffer, int offset, int count)
+    {
+        // If debugging, the raw traffic will be visible in the debug console
+        if (Debugger.IsAttached)
+        {
+            string data = Encoding.UTF8.GetString(buffer, offset, count);
+            Debug.WriteLine($"{header} {data}");
+        }
+    }
+}

test/PowerShellEditorServices.Test.E2E/Hosts/IAsyncLanguageServerHost.cs

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+#nullable enable
+
+using System;
+using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace PowerShellEditorServices.Test.E2E;
+
+/// <summary>
+/// Represents a debug adapter server host that can be started and stopped and provides streams for communication.
+/// </summary>
+public interface IAsyncLanguageServerHost : IAsyncDisposable
+{
+    // Start the host and return when the host is ready to communicate. It should return a tuple of a stream Reader and stream Writer for communication with the LSP. The underlying streams can be retrieved via baseStream propertyif needed.
+    Task<(StreamReader, StreamWriter)> Start(CancellationToken token = default);
+    // Stops the host and returns when the host has fully stopped. It should be idempotent, such that if called while the host is already stopping/stopped, it will have the same result
+    Task<bool> Stop(CancellationToken token = default);
+
+    // Optional to implement if more is required than a simple stop
+    async ValueTask IAsyncDisposable.DisposeAsync() => await Stop();
+}
+

test/PowerShellEditorServices.Test.E2E/DebugAdapterClientExtensions.cs → test/PowerShellEditorServices.Test.E2E/Hosts/IDebugAdapterClientExtensions.cs

@@ -4,14 +4,14 @@
 using System;
 using System.Threading.Tasks;
 using Microsoft.PowerShell.EditorServices.Handlers;
-using OmniSharp.Extensions.DebugAdapter.Client;
+using OmniSharp.Extensions.DebugAdapter.Protocol.Client;
 using OmniSharp.Extensions.DebugAdapter.Protocol.Requests;
 
 namespace PowerShellEditorServices.Test.E2E
 {
-    public static class DebugAdapterClientExtensions
+    public static class IDebugAdapterClientExtensions
     {
-        public static async Task LaunchScript(this DebugAdapterClient debugAdapterClient, string script, TaskCompletionSource<object> started, string executeMode = "DotSource")
+        public static async Task LaunchScript(this IDebugAdapterClient debugAdapterClient, string script, string executeMode = "DotSource")
         {
             _ = await debugAdapterClient.Launch(
                 new PsesLaunchRequestArguments
@@ -22,9 +22,6 @@ public static async Task LaunchScript(this DebugAdapterClient debugAdapterClient
                     CreateTemporaryIntegratedConsole = false,
                     ExecuteMode = executeMode,
                 }) ?? throw new Exception("Launch response was null.");
-
-            // This will check to see if we received the Initialized event from the server.
-            await started.Task;
         }
     }
 }

test/PowerShellEditorServices.Test.E2E/Hosts/PsesStdioLanguageServerProcessHost.cs

@@ -0,0 +1,85 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Reflection;
+
+namespace PowerShellEditorServices.Test.E2E;
+
+/// <summary>
+/// A <see cref="ServerManager"/> is responsible for launching or attaching to a language server, providing access to its input and output streams, and tracking its lifetime.
+/// </summary>
+internal class PsesStdioLanguageServerProcessHost(bool isDebugAdapter)
+: StdioLanguageServerProcessHost(PwshExe, GeneratePsesArguments(isDebugAdapter))
+{
+    protected static readonly string s_binDir =
+        Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location);
+
+    private static readonly string s_bundledModulePath = new FileInfo(Path.Combine(
+        s_binDir, "..", "..", "..", "..", "..", "module")).FullName;
+
+    private static readonly string s_sessionDetailsPath = Path.Combine(
+        s_binDir, $"pses_test_sessiondetails_{Path.GetRandomFileName()}");
+
+    private static readonly string s_logPath = Path.Combine(
+        s_binDir, $"pses_test_logs_{Path.GetRandomFileName()}");
+
+    private const string s_logLevel = "Diagnostic";
+    private static readonly string[] s_featureFlags = { "PSReadLine" };
+    private const string s_hostName = "TestHost";
+    private const string s_hostProfileId = "TestHost";
+    private const string s_hostVersion = "1.0.0";
+
+    // Adjust the environment variable if wanting to test with 5.1 or a specific pwsh path
+    public static string PwshExe { get; } = Environment.GetEnvironmentVariable("PWSH_EXE_NAME") ?? "pwsh";
+    public static bool IsWindowsPowerShell { get; } = PwshExe.EndsWith("powershell");
+    public static bool RunningInConstrainedLanguageMode { get; } =
+        Environment.GetEnvironmentVariable("__PSLockdownPolicy", EnvironmentVariableTarget.Machine) != null;
+
+    private static string[] GeneratePsesArguments(bool isDebugAdapter)
+    {
+        List<string> args = new()
+            {
+                "&",
+                SingleQuoteEscape(Path.Combine(s_bundledModulePath, "PowerShellEditorServices", "Start-EditorServices.ps1")),
+                "-LogPath",
+                SingleQuoteEscape(s_logPath),
+                "-LogLevel",
+                s_logLevel,
+                "-SessionDetailsPath",
+                SingleQuoteEscape(s_sessionDetailsPath),
+                "-FeatureFlags",
+                string.Join(',', s_featureFlags),
+                "-HostName",
+                s_hostName,
+                "-HostProfileId",
+                s_hostProfileId,
+                "-HostVersion",
+                s_hostVersion,
+                "-BundledModulesPath",
+                SingleQuoteEscape(s_bundledModulePath),
+                "-Stdio"
+            };
+
+        if (isDebugAdapter)
+        {
+            args.Add("-DebugServiceOnly");
+        }
+
+        string base64Str = Convert.ToBase64String(
+            System.Text.Encoding.Unicode.GetBytes(string.Join(' ', args)));
+
+        return
+        [
+            "-NoLogo",
+                "-NoProfile",
+                "-EncodedCommand",
+                base64Str
+        ];
+    }
+
+    private static string SingleQuoteEscape(string str) => $"'{str.Replace("'", "''")}'";
+}
+

test/PowerShellEditorServices.Test.E2E/Hosts/StdioLanguageServerProcessHost.cs

@@ -0,0 +1,116 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+#nullable enable
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace PowerShellEditorServices.Test.E2E;
+
+/// <summary>
+/// Hosts a language server process that communicates over stdio
+/// </summary>
+internal class StdioLanguageServerProcessHost(string fileName, IEnumerable<string> argumentList) : IAsyncLanguageServerHost
+{
+    // The PSES process that will be started and managed
+    private readonly Process process = new()
+    {
+        EnableRaisingEvents = true,
+        StartInfo = new ProcessStartInfo(fileName, argumentList)
+        {
+            RedirectStandardInput = true,
+            RedirectStandardOutput = true,
+            RedirectStandardError = true,
+            UseShellExecute = false,
+            CreateNoWindow = true
+        }
+    };
+
+    // Track the state of the startup
+    private TaskCompletionSource<(StreamReader, StreamWriter)>? startTcs;
+    private TaskCompletionSource<bool>? stopTcs;
+
+    // Starts the process. Returns when the process has started and streams are available.
+    public async Task<(StreamReader, StreamWriter)> Start(CancellationToken token = default)
+    {
+        // Runs this once upon process exit to clean up the state.
+        EventHandler? exitHandler = null;
+        exitHandler = (sender, e) =>
+        {
+            // Complete the stopTcs task when the process finally exits, allowing stop to complete
+            stopTcs?.TrySetResult(true);
+            stopTcs = null;
+            startTcs = null;
+            process.Exited -= exitHandler;
+        };
+        process.Exited += exitHandler;
+
+        if (stopTcs is not null)
+        {
+            throw new InvalidOperationException("The process is currently stopping and cannot be started.");
+        }
+
+        // Await the existing task if we have already started, making this operation idempotent
+        if (startTcs is not null)
+        {
+            return await startTcs.Task;
+        }
+
+        // Initiate a new startTcs to track the startup
+        startTcs = new();
+
+        token.ThrowIfCancellationRequested();
+
+        // Should throw if there are any startup problems such as invalid path, etc.
+        process.Start();
+
+        // According to the source the streams should be allocated synchronously after the process has started, however it's not super clear so we will put this here in case there is an explicit race condition.
+        if (process.StandardInput.BaseStream is null || process.StandardOutput.BaseStream is null)
+        {
+            throw new InvalidOperationException("The process has started but the StandardInput or StandardOutput streams are not available. This should never happen and is probably a race condition, please report it to PowerShellEditorServices.");
+        }
+
+        startTcs.SetResult((
+            process.StandardOutput,
+            process.StandardInput
+        ));
+
+        // Return the result of the completion task
+        return await startTcs.Task;
+    }
+
+    public async Task WaitForExit(CancellationToken token = default)
+    {
+        AssertStarting();
+        await process.WaitForExitAsync(token);
+    }
+
+    /// <summary>
+    /// Determines if the process is in the starting state and throws if not.
+    /// </summary>
+    private void AssertStarting()
+    {
+        if (startTcs is null)
+        {
+            throw new InvalidOperationException("The process is not starting/started, use Start() first.");
+        }
+    }
+
+    public async Task<bool> Stop(CancellationToken token = default)
+    {
+        AssertStarting();
+        if (stopTcs is not null)
+        {
+            return await stopTcs.Task;
+        }
+        stopTcs = new();
+        token.ThrowIfCancellationRequested();
+        process.Kill();
+        await process.WaitForExitAsync(token);
+        return true;
+    }
+}

test/PowerShellEditorServices.Test.E2E/LSPTestsFixures.cs → test/PowerShellEditorServices.Test.E2E/LSPTestsFixtures.cs

@@ -8,9 +8,9 @@
 using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.Extensions.Logging;
 using Microsoft.PowerShell.EditorServices.Logging;
 using Microsoft.PowerShell.EditorServices.Services.Configuration;
+using Nerdbank.Streams;
 using Newtonsoft.Json.Linq;
 using OmniSharp.Extensions.LanguageServer.Client;
 using OmniSharp.Extensions.LanguageServer.Protocol;
@@ -38,27 +38,30 @@ public class LSPTestsFixture : IAsyncLifetime
         internal List<PsesTelemetryEvent> TelemetryEvents = new();
         public ITestOutputHelper Output { get; set; }
 
-        protected PsesStdioProcess _psesProcess;
-        public int ProcessId => _psesProcess.Id;
+        internal PsesStdioLanguageServerProcessHost _psesHost = new(IsDebugAdapterTests);
 
         public async Task InitializeAsync()
         {
-            LoggerFactory factory = new();
-            _psesProcess = new PsesStdioProcess(factory, IsDebugAdapterTests);
-            await _psesProcess.Start();
+            (StreamReader stdout, StreamWriter stdin) = await _psesHost.Start();
+
+            // Splice the streams together and enable debug logging of all messages sent and received
+            DebugOutputStream psesStream = new(
+                FullDuplexStream.Splice(stdout.BaseStream, stdin.BaseStream)
+            );
 
             DirectoryInfo testDir =
                 Directory.CreateDirectory(Path.Combine(s_binDir, Path.GetRandomFileName()));
 
             PsesLanguageClient = LanguageClient.PreInit(options =>
             {
                 options
-                    .WithInput(_psesProcess.OutputStream)
-                    .WithOutput(_psesProcess.InputStream)
+                    .WithInput(psesStream)
+                    .WithOutput(psesStream)
                     .WithWorkspaceFolder(DocumentUri.FromFileSystemPath(testDir.FullName), "testdir")
                     .WithInitializationOptions(new { EnableProfileLoading = false })
                     .OnPublishDiagnostics(diagnosticParams => Diagnostics.AddRange(diagnosticParams.Diagnostics.Where(d => d != null)))
-                    .OnLogMessage(logMessageParams => {
+                    .OnLogMessage(logMessageParams =>
+                    {
                         Output?.WriteLine($"{logMessageParams.Type}: {logMessageParams.Message}");
                         Messages.Add(logMessageParams);
                     })
@@ -98,7 +101,7 @@ public async Task InitializeAsync()
         public async Task DisposeAsync()
         {
             await PsesLanguageClient.Shutdown();
-            await _psesProcess.Stop();
+            await _psesHost.Stop();
             PsesLanguageClient?.Dispose();
         }
     }