mxaccessgw/src/ZB.MOM.WW.MxGateway.Tests/Gateway/Sessions/SessionWorkerClientFactoryFakeWorkerTests.cs

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

namespace ZB.MOM.WW.MxGateway.Tests.Gateway.Sessions;

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

    // Anti-hang guard for the CreateAsync failure/timeout tests. It only exists to stop a genuine
    // hang from wedging the run, so it must be far larger than any in-test semantic timeout (the
    // factory startup timeout being exercised) — otherwise, under CI load, this WaitAsync net can
    // trip before the factory's own timeout propagates and the test observes .NET's generic
    // "The operation has timed out." instead of the asserted "did not complete startup".
    private static readonly TimeSpan HangGuardTimeout = TimeSpan.FromSeconds(30);

    private readonly List<IWorkerTaskLauncher> _launchers = [];

    /// <summary>
    /// Awaits every scripted worker task so an unhandled exception fails the owning test
    /// instead of surfacing later as an unobserved <see cref="TaskScheduler.UnobservedTaskException"/>.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public async ValueTask DisposeAsync()
    {
        foreach (IWorkerTaskLauncher launcher in _launchers)
        {
            await launcher.ObserveWorkerTaskAsync(TestTimeout);
        }
    }

    /// <summary>Verifies that the factory creates a ready worker client with a scripted fake worker process.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task CreateAsync_WithScriptedFakeWorker_ReturnsReadyClient()
    {
        ScriptedFakeWorkerProcessLauncher launcher = Track(new ScriptedFakeWorkerProcessLauncher());
        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);
    }

    /// <summary>Verifies that a failed fake worker startup throws a worker client exception.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task CreateAsync_WhenFakeWorkerStartupFails_ThrowsWorkerClientException()
    {
        FailingStartupWorkerProcessLauncher launcher = Track(new FailingStartupWorkerProcessLauncher());
        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(HangGuardTimeout));

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

    /// <summary>Verifies that a worker that never sends ready times out and is killed.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task CreateAsync_WhenFakeWorkerNeverSendsReady_TimesOutAndKillsWorker()
    {
        NeverReadyWorkerProcessLauncher launcher = Track(new NeverReadyWorkerProcessLauncher());
        using GatewayMetrics metrics = new();
        SessionWorkerClientFactory factory = new(
            launcher,
            Options.Create(CreateOptions(startupTimeoutSeconds: 1)),
            metrics,
            NullLoggerFactory.Instance);
        GatewaySession session = CreateSession(startupTimeout: TimeSpan.FromSeconds(1));

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

        Assert.Contains("did not complete startup", exception.Message);
        Assert.Equal(1, launcher.Process.KillCount);
        Assert.True(launcher.Process.IsDisposed);
    }

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

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

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

    private T Track<T>(T launcher)
        where T : IWorkerTaskLauncher
    {
        _launchers.Add(launcher);

        return launcher;
    }

    /// <summary>
    /// A fake worker launcher that runs a scripted worker on a background task and exposes
    /// that task so the owning test observes it rather than leaking an unobserved fault.
    /// </summary>
    private interface IWorkerTaskLauncher : IWorkerProcessLauncher
    {
        /// <summary>
        /// Awaits the scripted worker task within the timeout, swallowing only the pipe
        /// teardown faults expected when the worker client kills or disposes the worker.
        /// </summary>
        /// <param name="timeout">Maximum time to wait for the worker task.</param>
        /// <returns>A task that represents the asynchronous operation.</returns>
        Task ObserveWorkerTaskAsync(TimeSpan timeout);
    }

    /// <summary>
    /// Awaits a scripted worker task, treating cancellation and pipe-disconnect I/O faults as
    /// the expected outcome of the worker client tearing the worker down, and rethrowing anything else.
    /// </summary>
    private static async Task ObserveWorkerTaskAsync(Task workerTask, TimeSpan timeout)
    {
        try
        {
            await workerTask.WaitAsync(timeout).ConfigureAwait(false);
        }
        catch (OperationCanceledException)
        {
            // Expected: the worker client cancelled the scripted worker during teardown.
        }
        catch (IOException)
        {
            // Expected: the gateway pipe was closed when the worker client disposed.
        }
    }

    /// <summary>Fake worker launcher that connects a scripted fake worker harness.</summary>
    private sealed class ScriptedFakeWorkerProcessLauncher : IWorkerTaskLauncher
    {
        /// <summary>The fake process ID used by the scripted launcher.</summary>
        public const int ProcessId = 2468;
        private readonly FakeWorkerProcess _process = new(ProcessId);

        /// <summary>Gets the connected fake worker harness.</summary>
        public FakeWorkerHarness? Harness { get; private set; }

        /// <summary>Gets the scripted worker task.</summary>
        public Task WorkerTask { get; private set; } = Task.CompletedTask;

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

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

        /// <inheritdoc />
        public Task ObserveWorkerTaskAsync(TimeSpan timeout) =>
            SessionWorkerClientFactoryFakeWorkerTests.ObserveWorkerTaskAsync(WorkerTask, timeout);

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

    /// <summary>Fake worker launcher that fails during startup with protocol version mismatch.</summary>
    private sealed class FailingStartupWorkerProcessLauncher : IWorkerTaskLauncher
    {
        /// <summary>Gets the fake worker process.</summary>
        public FakeWorkerProcess Process { get; } = new(processId: 3579);

        /// <summary>Gets the scripted worker task.</summary>
        public Task WorkerTask { get; private set; } = Task.CompletedTask;

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

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

        /// <inheritdoc />
        public Task ObserveWorkerTaskAsync(TimeSpan timeout) =>
            SessionWorkerClientFactoryFakeWorkerTests.ObserveWorkerTaskAsync(WorkerTask, timeout);

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

    /// <summary>Fake worker launcher that never completes startup, simulating a hung worker.</summary>
    private sealed class NeverReadyWorkerProcessLauncher : IWorkerTaskLauncher
    {
        private readonly CancellationTokenSource _stop = new();

        /// <summary>Gets the fake worker process.</summary>
        public FakeWorkerProcess Process { get; } = new(processId: 4680);

        /// <summary>Gets the scripted worker task.</summary>
        public Task WorkerTask { get; private set; } = Task.CompletedTask;

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

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

        /// <inheritdoc />
        public async Task ObserveWorkerTaskAsync(TimeSpan timeout)
        {
            // The scripted worker parks on an infinite delay; cancel it so disposal observes
            // the task instead of leaking it as an unobserved fault.
            await _stop.CancelAsync().ConfigureAwait(false);
            await SessionWorkerClientFactoryFakeWorkerTests
                .ObserveWorkerTaskAsync(WorkerTask, timeout)
                .ConfigureAwait(false);
            _stop.Dispose();
        }

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

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

}