mxaccessgw/src/MxGateway.Tests/Gateway/Workers/WorkerProcessLauncherTests.cs

using System.Diagnostics;
using Microsoft.Extensions.Options;
using MxGateway.Contracts;
using MxGateway.Server.Configuration;
using MxGateway.Server.Metrics;
using MxGateway.Server.Workers;

namespace MxGateway.Tests.Gateway.Workers;

public sealed class WorkerProcessLauncherTests
{
    private const string SessionId = "session-1";
    private const string PipeName = "mxaccess-gateway-123-session-1";
    private const string Nonce = "super-secret-nonce";

    [Fact]
    public async Task LaunchAsync_WithValidWorker_StartsProcessWithBootstrapArgumentsAndNonceEnvironment()
    {
        using TestDirectory directory = TestDirectory.Create();
        string executablePath = directory.CreateWorkerExecutable(machine: 0x014c);
        FakeWorkerProcess process = new(processId: 1234);
        FakePipeReservation pipeReservation = new();
        FakeWorkerProcessFactory processFactory = new(process);
        GatewayMetrics metrics = new();
        WorkerProcessLauncher launcher = CreateLauncher(executablePath, processFactory, new SucceedingStartupProbe(), metrics);

        using WorkerProcessHandle handle = await launcher.LaunchAsync(CreateRequest(pipeReservation));

        Assert.Equal(1234, handle.ProcessId);
        Assert.Same(process, handle.Process);
        Assert.NotNull(processFactory.LastStartInfo);
        Assert.Equal(Path.GetFullPath(executablePath), processFactory.LastStartInfo.FileName);
        Assert.False(processFactory.LastStartInfo.UseShellExecute);
        Assert.True(processFactory.LastStartInfo.CreateNoWindow);
        Assert.Equal(
            ["--session-id", SessionId, "--pipe-name", PipeName, "--protocol-version", "1"],
            processFactory.LastStartInfo.ArgumentList);
        Assert.Equal(Nonce, processFactory.LastStartInfo.Environment[WorkerProcessLauncher.WorkerNonceEnvironmentVariableName]);
        Assert.DoesNotContain(Nonce, handle.CommandLine.ToString(), StringComparison.Ordinal);
        Assert.DoesNotContain(Nonce, string.Join(" ", handle.CommandLine.Arguments), StringComparison.Ordinal);
        Assert.False(pipeReservation.DisposeCalled);
        Assert.Equal(1, metrics.GetSnapshot().WorkersRunning);
    }

    [Fact]
    public async Task LaunchAsync_WhenStartupProbeFails_KillsAndDisposesWorker()
    {
        using TestDirectory directory = TestDirectory.Create();
        string executablePath = directory.CreateWorkerExecutable(machine: 0x014c);
        FakeWorkerProcess process = new(processId: 1234);
        FakePipeReservation pipeReservation = new();
        GatewayMetrics metrics = new();
        WorkerProcessLauncher launcher = CreateLauncher(
            executablePath,
            new FakeWorkerProcessFactory(process),
            new FailingStartupProbe(),
            metrics);

        WorkerProcessLaunchException exception =
            await Assert.ThrowsAsync<WorkerProcessLaunchException>(
                async () => await launcher.LaunchAsync(CreateRequest(pipeReservation)));

        Assert.Equal(WorkerProcessLaunchErrorCode.StartupFailed, exception.ErrorCode);
        Assert.True(process.KillCalled);
        Assert.True(process.DisposeCalled);
        Assert.True(pipeReservation.DisposeCalled);
        Assert.Equal(1, metrics.GetSnapshot().WorkerKills);
    }

    [Fact]
    public async Task LaunchAsync_WhenStartupTimesOut_KillsAndDisposesWorker()
    {
        using TestDirectory directory = TestDirectory.Create();
        string executablePath = directory.CreateWorkerExecutable(machine: 0x014c);
        FakeWorkerProcess process = new(processId: 1234);
        GatewayMetrics metrics = new();
        WorkerProcessLauncher launcher = CreateLauncher(
            executablePath,
            new FakeWorkerProcessFactory(process),
            new WaitingStartupProbe(),
            metrics,
            startupTimeoutSeconds: 1);

        WorkerProcessLaunchException exception =
            await Assert.ThrowsAsync<WorkerProcessLaunchException>(
                async () => await launcher.LaunchAsync(CreateRequest()));

        Assert.Equal(WorkerProcessLaunchErrorCode.StartupTimeout, exception.ErrorCode);
        Assert.True(process.KillCalled);
        Assert.True(process.DisposeCalled);
        Assert.Equal(1, metrics.GetSnapshot().WorkerKills);
    }

    [Fact]
    public async Task LaunchAsync_WhenExecutableDoesNotExist_FailsBeforeStartingProcess()
    {
        using TestDirectory directory = TestDirectory.Create();
        string executablePath = Path.Combine(directory.Path, "missing-worker.exe");
        FakeWorkerProcessFactory processFactory = new(new FakeWorkerProcess(processId: 1234));
        WorkerProcessLauncher launcher = CreateLauncher(executablePath, processFactory, new SucceedingStartupProbe());

        WorkerProcessLaunchException exception =
            await Assert.ThrowsAsync<WorkerProcessLaunchException>(
                async () => await launcher.LaunchAsync(CreateRequest()));

        Assert.Equal(WorkerProcessLaunchErrorCode.ExecutableNotFound, exception.ErrorCode);
        Assert.Null(processFactory.LastStartInfo);
    }

    [Fact]
    public async Task LaunchAsync_WhenExecutableArchitectureDoesNotMatch_FailsBeforeStartingProcess()
    {
        using TestDirectory directory = TestDirectory.Create();
        string executablePath = directory.CreateWorkerExecutable(machine: 0x8664);
        FakeWorkerProcessFactory processFactory = new(new FakeWorkerProcess(processId: 1234));
        WorkerProcessLauncher launcher = CreateLauncher(executablePath, processFactory, new SucceedingStartupProbe());

        WorkerProcessLaunchException exception =
            await Assert.ThrowsAsync<WorkerProcessLaunchException>(
                async () => await launcher.LaunchAsync(CreateRequest()));

        Assert.Equal(WorkerProcessLaunchErrorCode.InvalidExecutable, exception.ErrorCode);
        Assert.Null(processFactory.LastStartInfo);
    }

    [Fact]
    public async Task LaunchAsync_WhenWorkerAlreadyExited_FailsAndDisposesWorkerWithoutKill()
    {
        using TestDirectory directory = TestDirectory.Create();
        string executablePath = directory.CreateWorkerExecutable(machine: 0x014c);
        FakeWorkerProcess process = new(processId: 1234)
        {
            HasExited = true,
            ExitCode = 42,
        };
        WorkerProcessLauncher launcher = CreateLauncher(
            executablePath,
            new FakeWorkerProcessFactory(process),
            new WorkerProcessStartedProbe());

        WorkerProcessLaunchException exception =
            await Assert.ThrowsAsync<WorkerProcessLaunchException>(
                async () => await launcher.LaunchAsync(CreateRequest()));

        Assert.Equal(WorkerProcessLaunchErrorCode.StartupFailed, exception.ErrorCode);
        Assert.False(process.KillCalled);
        Assert.True(process.DisposeCalled);
    }

    private static WorkerProcessLauncher CreateLauncher(
        string executablePath,
        IWorkerProcessFactory processFactory,
        IWorkerStartupProbe startupProbe,
        GatewayMetrics? metrics = null,
        int startupTimeoutSeconds = 30)
    {
        GatewayOptions options = new()
        {
            Worker = new WorkerOptions
            {
                ExecutablePath = executablePath,
                RequiredArchitecture = WorkerArchitecture.X86,
                StartupTimeoutSeconds = startupTimeoutSeconds,
            },
        };

        return new WorkerProcessLauncher(
            Options.Create(options),
            processFactory,
            startupProbe,
            metrics ?? new GatewayMetrics());
    }

    private static WorkerProcessLaunchRequest CreateRequest(IDisposable? pipeReservation = null)
    {
        return new WorkerProcessLaunchRequest(
            SessionId,
            PipeName,
            GatewayContractInfo.WorkerProtocolVersion,
            Nonce,
            pipeReservation);
    }

    private sealed class FakeWorkerProcessFactory(IWorkerProcess process) : IWorkerProcessFactory
    {
        public ProcessStartInfo? LastStartInfo { get; private set; }

        public IWorkerProcess Start(ProcessStartInfo startInfo)
        {
            LastStartInfo = startInfo;
            return process;
        }
    }

    private sealed class FakeWorkerProcess(int processId) : IWorkerProcess
    {
        public int Id { get; } = processId;

        public bool HasExited { get; set; }

        public int? ExitCode { get; set; }

        public bool DisposeCalled { get; private set; }

        public bool KillCalled { get; private set; }

        public ValueTask WaitForExitAsync(CancellationToken cancellationToken)
        {
            return ValueTask.CompletedTask;
        }

        public void Kill(bool entireProcessTree)
        {
            Assert.True(entireProcessTree);
            KillCalled = true;
            HasExited = true;
        }

        public void Dispose()
        {
            DisposeCalled = true;
        }
    }

    private sealed class SucceedingStartupProbe : IWorkerStartupProbe
    {
        public Task WaitUntilReadyAsync(
            IWorkerProcess process,
            WorkerProcessLaunchRequest request,
            CancellationToken cancellationToken)
        {
            return Task.CompletedTask;
        }
    }

    private sealed class FailingStartupProbe : IWorkerStartupProbe
    {
        public Task WaitUntilReadyAsync(
            IWorkerProcess process,
            WorkerProcessLaunchRequest request,
            CancellationToken cancellationToken)
        {
            throw new InvalidOperationException("Fake worker startup failed.");
        }
    }

    private sealed class WaitingStartupProbe : IWorkerStartupProbe
    {
        public async Task WaitUntilReadyAsync(
            IWorkerProcess process,
            WorkerProcessLaunchRequest request,
            CancellationToken cancellationToken)
        {
            await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken);
        }
    }

    private sealed class FakePipeReservation : IDisposable
    {
        public bool DisposeCalled { get; private set; }

        public void Dispose()
        {
            DisposeCalled = true;
        }
    }

    private sealed class TestDirectory : IDisposable
    {
        private TestDirectory(string path)
        {
            Path = path;
        }

        public string Path { get; }

        public static TestDirectory Create()
        {
            string path = System.IO.Path.Combine(System.IO.Path.GetTempPath(), $"mxgateway-tests-{Guid.NewGuid():N}");
            Directory.CreateDirectory(path);

            return new TestDirectory(path);
        }

        public string CreateWorkerExecutable(ushort machine)
        {
            string path = System.IO.Path.Combine(Path, "MxGateway.Worker.exe");
            byte[] bytes = new byte[0x100];
            bytes[0] = (byte)'M';
            bytes[1] = (byte)'Z';
            BitConverter.GetBytes(0x80).CopyTo(bytes, 0x3c);
            bytes[0x80] = (byte)'P';
            bytes[0x81] = (byte)'E';
            bytes[0x82] = 0;
            bytes[0x83] = 0;
            BitConverter.GetBytes(machine).CopyTo(bytes, 0x84);
            File.WriteAllBytes(path, bytes);

            return path;
        }

        public void Dispose()
        {
            Directory.Delete(Path, recursive: true);
        }
    }
}