mxaccessgw/src/MxGateway.Tests/Gateway/Sessions/SessionWorkerClientFactoryFakeWorkerTests.cs

using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using MxGateway.Contracts;
using MxGateway.Contracts.Proto;
using MxGateway.Server.Configuration;
using MxGateway.Server.Metrics;
using MxGateway.Server.Sessions;
using MxGateway.Server.Workers;
using MxGateway.Tests.Gateway.Workers.Fakes;

namespace MxGateway.Tests.Gateway.Sessions;

public sealed class SessionWorkerClientFactoryFakeWorkerTests
{
    private static readonly TimeSpan TestTimeout = TimeSpan.FromSeconds(5);

    [Fact]
    public async Task CreateAsync_WithScriptedFakeWorker_ReturnsReadyClient()
    {
        ScriptedFakeWorkerProcessLauncher launcher = new();
        using GatewayMetrics metrics = new();
        SessionWorkerClientFactory factory = new(
            launcher,
            Options.Create(CreateOptions()),
            metrics,
            NullLoggerFactory.Instance);
        GatewaySession session = CreateSession();

        await using IWorkerClient workerClient = await factory.CreateAsync(
            session,
            CancellationToken.None);

        Assert.Equal(WorkerClientState.Ready, workerClient.State);
        Assert.Equal(ScriptedFakeWorkerProcessLauncher.ProcessId, workerClient.ProcessId);
        Assert.NotNull(launcher.Harness);

        Task<WorkerCommandReply> invokeTask = workerClient.InvokeAsync(
            CreateCommand(MxCommandKind.Ping),
            TestTimeout,
            CancellationToken.None);
        WorkerEnvelope commandEnvelope = await launcher.Harness.ReadCommandAsync();
        await launcher.Harness.ReplyToCommandAsync(commandEnvelope);
        WorkerCommandReply reply = await invokeTask.WaitAsync(TestTimeout);

        Assert.Equal(MxCommandKind.Ping, reply.Reply.Kind);
        Assert.Equal(ProtocolStatusCode.Ok, reply.Reply.ProtocolStatus.Code);
    }

    [Fact]
    public async Task CreateAsync_WhenFakeWorkerStartupFails_ThrowsWorkerClientException()
    {
        FailingStartupWorkerProcessLauncher launcher = new();
        using GatewayMetrics metrics = new();
        SessionWorkerClientFactory factory = new(
            launcher,
            Options.Create(CreateOptions()),
            metrics,
            NullLoggerFactory.Instance);
        GatewaySession session = CreateSession();

        WorkerClientException exception = await Assert.ThrowsAsync<WorkerClientException>(
            async () => await factory.CreateAsync(session, CancellationToken.None).WaitAsync(TestTimeout));

        Assert.Equal(WorkerClientErrorCode.ProtocolViolation, exception.ErrorCode);
        Assert.True(launcher.Process.IsDisposed);
    }

    private static GatewayOptions CreateOptions()
    {
        return new GatewayOptions
        {
            Worker = new WorkerOptions
            {
                StartupTimeoutSeconds = 5,
                ShutdownTimeoutSeconds = 5,
                HeartbeatIntervalSeconds = 30,
                HeartbeatGraceSeconds = 30,
                MaxMessageBytes = WorkerFrameProtocolOptions.DefaultMaxMessageBytes,
            },
            Events = new EventOptions
            {
                QueueCapacity = 16,
            },
        };
    }

    private static GatewaySession CreateSession()
    {
        return new GatewaySession(
            FakeWorkerHarness.DefaultSessionId,
            GatewayContractInfo.DefaultBackendName,
            $"mxaccessgw-session-fake-worker-{Guid.NewGuid():N}",
            FakeWorkerHarness.DefaultNonce,
            "test-client",
            "fake-worker-session-test",
            "client-correlation-1",
            TestTimeout,
            TestTimeout,
            TestTimeout,
            DateTimeOffset.UtcNow);
    }

    private static WorkerCommand CreateCommand(MxCommandKind kind)
    {
        return new WorkerCommand
        {
            Command = new MxCommand
            {
                Kind = kind,
            },
        };
    }

    private sealed class ScriptedFakeWorkerProcessLauncher : IWorkerProcessLauncher
    {
        public const int ProcessId = 2468;
        private readonly FakeWorkerProcess _process = new(ProcessId);

        public FakeWorkerHarness? Harness { get; private set; }

        public Task<WorkerProcessHandle> LaunchAsync(
            WorkerProcessLaunchRequest request,
            CancellationToken cancellationToken = default)
        {
            _ = RunWorkerAsync(request, cancellationToken);

            return Task.FromResult(CreateHandle(_process));
        }

        private async Task RunWorkerAsync(
            WorkerProcessLaunchRequest request,
            CancellationToken cancellationToken)
        {
            Harness = await FakeWorkerHarness.ConnectToGatewayPipeAsync(
                request.SessionId,
                request.Nonce,
                request.PipeName,
                request.ProtocolVersion,
                cancellationToken: cancellationToken).ConfigureAwait(false);
            await Harness.CompleteStartupAsync(ProcessId, cancellationToken: cancellationToken).ConfigureAwait(false);
        }
    }

    private sealed class FailingStartupWorkerProcessLauncher : IWorkerProcessLauncher
    {
        public FakeWorkerProcess Process { get; } = new(processId: 3579);

        public Task<WorkerProcessHandle> LaunchAsync(
            WorkerProcessLaunchRequest request,
            CancellationToken cancellationToken = default)
        {
            _ = RunWorkerAsync(request, cancellationToken);

            return Task.FromResult(CreateHandle(Process));
        }

        private async Task RunWorkerAsync(
            WorkerProcessLaunchRequest request,
            CancellationToken cancellationToken)
        {
            await using FakeWorkerHarness harness = await FakeWorkerHarness.ConnectToGatewayPipeAsync(
                request.SessionId,
                request.Nonce,
                request.PipeName,
                request.ProtocolVersion,
                cancellationToken: cancellationToken).ConfigureAwait(false);
            _ = await harness.ReadGatewayEnvelopeAsync(cancellationToken).ConfigureAwait(false);
            await harness.SendWorkerHelloAsync(
                workerProcessId: Process.Id,
                workerProtocolVersion: request.ProtocolVersion + 1,
                cancellationToken: cancellationToken).ConfigureAwait(false);
        }
    }

    private static WorkerProcessHandle CreateHandle(IWorkerProcess process)
    {
        return new WorkerProcessHandle(
            process,
            new WorkerProcessCommandLine("fake-worker.exe", []),
            DateTimeOffset.UtcNow);
    }

    private sealed class FakeWorkerProcess(int processId) : IWorkerProcess
    {
        private bool _disposed;

        public int Id { get; } = processId;

        public bool HasExited { get; private set; }

        public int? ExitCode { get; private set; }

        public int KillCount { get; private set; }

        public ValueTask WaitForExitAsync(CancellationToken cancellationToken)
        {
            HasExited = true;
            ExitCode = 0;
            return ValueTask.CompletedTask;
        }

        public void Kill(bool entireProcessTree)
        {
            KillCount++;
            HasExited = true;
            ExitCode = -1;
        }

        public void Dispose()
        {
            _disposed = true;
        }

        public bool IsDisposed => _disposed;
    }
}