mxaccessgw/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Workers/WorkerProcessLauncherTests.cs

using System.Diagnostics;
using Microsoft.Extensions.Options;
using ZB.MOM.WW.MxGateway.Contracts;
using ZB.MOM.WW.MxGateway.Server.Configuration;
using ZB.MOM.WW.MxGateway.Server.Metrics;
using ZB.MOM.WW.MxGateway.Server.Workers;
using ZB.MOM.WW.MxGateway.Tests.TestSupport;

namespace ZB.MOM.WW.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";

    /// <summary>Verifies that a valid worker executable starts with correct bootstrap arguments and nonce environment variable.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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.Equal(
            "2000",
            processFactory.LastStartInfo.Environment[
                WorkerProcessLauncher.WorkerPipeConnectAttemptTimeoutEnvironmentVariableName]);
        Assert.DoesNotContain(Nonce, handle.CommandLine.ToString(), StringComparison.Ordinal);
        Assert.DoesNotContain(Nonce, string.Join(" ", handle.CommandLine.Arguments), StringComparison.Ordinal);
        Assert.False(pipeReservation.DisposeCalled);
        Assert.Equal(0, metrics.GetSnapshot().WorkersRunning);
    }

    /// <summary>Verifies that a failed startup probe kills and disposes the worker process.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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.KillEntireProcessTree);
        Assert.True(process.DisposeCalled);
        Assert.True(pipeReservation.DisposeCalled);
        Assert.Equal(1, metrics.GetSnapshot().WorkerKills);
    }

    /// <summary>Verifies that transient startup probe failures are retried without respawning the worker process.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task LaunchAsync_WhenStartupProbeFailsTransiently_RetriesWithoutRespawningWorker()
    {
        using TestDirectory directory = TestDirectory.Create();
        string executablePath = directory.CreateWorkerExecutable(machine: 0x014c);
        FakeWorkerProcess process = new(processId: 1234);
        FakeWorkerProcessFactory processFactory = new(process);
        GatewayMetrics metrics = new();
        WorkerProcessLauncher launcher = CreateLauncher(
            executablePath,
            processFactory,
            new TransientStartupProbe(failuresBeforeSuccess: 1),
            metrics,
            startupProbeRetryAttempts: 2,
            startupProbeRetryDelayMilliseconds: 1);

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

        Assert.Same(process, handle.Process);
        Assert.Equal(1, processFactory.StartCount);
        Assert.False(process.KillCalled);
        GatewayMetricsSnapshot snapshot = metrics.GetSnapshot();
        Assert.Equal(1, snapshot.RetryAttempts);
        Assert.Equal(1, snapshot.RetryAttemptsByArea["worker_startup"]);
    }

    /// <summary>Verifies that a startup probe timeout kills and disposes the worker process.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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.KillEntireProcessTree);
        Assert.True(process.DisposeCalled);
        Assert.Equal(1, metrics.GetSnapshot().WorkerKills);
    }

    /// <summary>Verifies that a missing worker executable fails before attempting to start the process.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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);
    }

    /// <summary>Verifies that a worker executable with mismatched architecture fails before attempting to start.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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);
    }

    /// <summary>Verifies that a worker that has already exited fails and disposes without additional killing.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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,
        int startupProbeRetryAttempts = 3,
        int startupProbeRetryDelayMilliseconds = 250)
    {
        GatewayOptions options = new()
        {
            Worker = new WorkerOptions
            {
                ExecutablePath = executablePath,
                RequiredArchitecture = WorkerArchitecture.X86,
                StartupTimeoutSeconds = startupTimeoutSeconds,
                StartupProbeRetryAttempts = startupProbeRetryAttempts,
                StartupProbeRetryDelayMilliseconds = startupProbeRetryDelayMilliseconds,
            },
        };

        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);
    }

    /// <summary>Fake worker process factory for testing process launch logic.</summary>
    private sealed class FakeWorkerProcessFactory(IWorkerProcess process) : IWorkerProcessFactory
    {
        /// <summary>Gets the most recent process start information.</summary>
        public ProcessStartInfo? LastStartInfo { get; private set; }

        /// <summary>Gets the number of times the process factory has started a process.</summary>
        public int StartCount { get; private set; }

        /// <inheritdoc />
        public IWorkerProcess Start(ProcessStartInfo startInfo)
        {
            StartCount++;
            LastStartInfo = startInfo;
            return process;
        }
    }

    /// <summary>Fake startup probe that immediately succeeds.</summary>
    private sealed class SucceedingStartupProbe : IWorkerStartupProbe
    {
        /// <inheritdoc />
        public Task WaitUntilReadyAsync(
            IWorkerProcess process,
            WorkerProcessLaunchRequest request,
            CancellationToken cancellationToken)
        {
            return Task.CompletedTask;
        }
    }

    /// <summary>Fake startup probe that always fails.</summary>
    private sealed class FailingStartupProbe : IWorkerStartupProbe
    {
        /// <inheritdoc />
        public Task WaitUntilReadyAsync(
            IWorkerProcess process,
            WorkerProcessLaunchRequest request,
            CancellationToken cancellationToken)
        {
            throw new InvalidOperationException("Fake worker startup failed.");
        }
    }

    /// <summary>Fake startup probe that waits indefinitely to simulate a startup timeout.</summary>
    private sealed class WaitingStartupProbe : IWorkerStartupProbe
    {
        /// <inheritdoc />
        public async Task WaitUntilReadyAsync(
            IWorkerProcess process,
            WorkerProcessLaunchRequest request,
            CancellationToken cancellationToken)
        {
            await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken);
        }
    }

    /// <summary>Fake startup probe that fails a configurable number of times before succeeding.</summary>
    private sealed class TransientStartupProbe(int failuresBeforeSuccess) : IWorkerStartupProbe
    {
        private int _attempts;

        /// <inheritdoc />
        public Task WaitUntilReadyAsync(
            IWorkerProcess process,
            WorkerProcessLaunchRequest request,
            CancellationToken cancellationToken)
        {
            if (Interlocked.Increment(ref _attempts) <= failuresBeforeSuccess)
            {
                throw new IOException("The worker pipe was not ready yet.");
            }

            return Task.CompletedTask;
        }
    }

    /// <summary>Fake pipe reservation for testing pipe lifecycle.</summary>
    private sealed class FakePipeReservation : IDisposable
    {
        /// <summary>Gets a value indicating whether the Dispose method was called.</summary>
        public bool DisposeCalled { get; private set; }

        /// <summary>Records that the pipe reservation was released.</summary>
        public void Dispose()
        {
            DisposeCalled = true;
        }
    }

    /// <summary>Test helper that creates and cleans up a temporary directory for worker executable tests.</summary>
    private sealed class TestDirectory : IDisposable
    {
        private TestDirectory(string path)
        {
            Path = path;
        }

        /// <summary>Gets the path to the temporary test directory.</summary>
        public string Path { get; }

        /// <summary>Creates a new temporary directory for testing.</summary>
        /// <returns>The created <see cref="TestDirectory"/> wrapper.</returns>
        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);
        }

        /// <summary>Creates a fake PE executable with the specified machine architecture for testing.</summary>
        /// <param name="machine">PE machine type constant (0x014c for x86, 0x8664 for x64).</param>
        /// <returns>Full path to the created executable file.</returns>
        public string CreateWorkerExecutable(ushort machine)
        {
            string path = System.IO.Path.Combine(Path, "ZB.MOM.WW.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;
        }

        /// <summary>Deletes the temporary test directory and its contents.</summary>
        public void Dispose()
        {
            Directory.Delete(Path, recursive: true);
        }
    }
}