ScadaBridge/tests/ZB.MOM.WW.ScadaBridge.Communication.Tests/Grpc/SitePairChannelProviderDisposalTests.cs

using Grpc.Net.Client;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using ZB.MOM.WW.ScadaBridge.Communication.Grpc;

namespace ZB.MOM.WW.ScadaBridge.Communication.Tests.Grpc;

/// <summary>
/// A channel replaced by the 60 s DB address refresh (or an admin
/// edit) must not be disposed out from under an in-flight <see cref="SitePairChannelProvider.ExecuteAsync{T}"/>
/// call. <c>ExecuteAsync</c> reads the channel under the pair gate but runs the RPC outside it, so an
/// inline dispose aborted whatever was running — a multi-minute <c>DeployArtifactsCommand</c>
/// included — with a status the cross-node retry deliberately never engages on.
/// </summary>
public class SitePairChannelProviderDisposalTests
{
    private static SitePairChannelProvider Build() => new(
        new NoKeyProvider(),
        Options.Create(new CommunicationOptions()),
        NullLogger<SitePairChannelProvider>.Instance);

    /// <summary>Verifies that shutdown retains the active channel and disposes it only once the call completes.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task Shutdown_RetainsActiveChannelThenDisposesItAtCompletion()
    {
        using var provider = Build();
        provider.UpdateSite("site-1", "http://localhost:15101", null);
        var entered = new TaskCompletionSource<GrpcChannel>(TaskCreationOptions.RunContinuationsAsynchronously);
        var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        var call = provider.ExecuteAsync("site-1", async (channel, ct) =>
        {
            entered.SetResult(channel);
            await release.Task;
            _ = channel.CreateCallInvoker();
            return 42;
        }, CancellationToken.None);
        var channel = await entered.Task;
        try
        {
            provider.Dispose();
            Assert.Equal(1, provider.ParkedChannelCount);
            _ = channel.CreateCallInvoker();
            Assert.Throws<ObjectDisposedException>(() => provider.UpdateSite("site-2", "http://localhost:15102", null));
        }
        finally { release.TrySetResult(); }
        Assert.Equal(42, await call);
        Assert.Equal(0, provider.ParkedChannelCount);
        Assert.Throws<ObjectDisposedException>(() => channel.CreateCallInvoker());
    }

    /// <summary>Verifies that a sweep keeps a retired channel alive until the in-flight call actually completes.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task RetiredChannel_SweepKeepsItAliveUntilActualCallCompletion()
    {
        using var provider = Build();
        provider.UpdateSite("site-1", "http://localhost:15101", null);
        var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        GrpcChannel? held = null;
        var call = provider.ExecuteAsync("site-1", async (channel, ct) =>
        {
            held = channel;
            entered.SetResult();
            await release.Task;
            return 42;
        }, CancellationToken.None);
        try
        {
            await entered.Task;
            provider.UpdateSite("site-1", "http://localhost:15102", null);
            // Simulates any elapsed grace, including waits longer than five minutes.
            provider.SweepGraveyard(force: true);
            Assert.Equal(1, provider.ParkedChannelCount);
            _ = held!.CreateCallInvoker();
        }
        finally { release.TrySetResult(); await call; }
        provider.SweepGraveyard(force: true);
        Assert.Equal(0, provider.ParkedChannelCount);
        Assert.Throws<ObjectDisposedException>(() => held!.CreateCallInvoker());
    }

    /// <summary>Verifies that an operation that may have been applied before the channel became unavailable is not replayed.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ApplyThenUnavailable_DoesNotReplayAnAmbiguousOperation()
    {
        using var provider = Build();
        provider.UpdateSite("site-1", "http://localhost:15101", "http://localhost:15102");
        var effects = 0;
        var error = await Record.ExceptionAsync(() => provider.ExecuteAsync("site-1", (channel, ct) =>
        {
            effects++;
            if (effects == 1)
                throw new global::Grpc.Core.RpcException(new global::Grpc.Core.Status(global::Grpc.Core.StatusCode.Unavailable, "response connection lost after apply"));
            return Task.FromResult(42);
        }, CancellationToken.None));
        Assert.Equal(1, effects);
        Assert.NotNull(error);
    }

    /// <summary>Verifies that a secondary node that applies the operation but loses the response reports an unknown outcome without another replay.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task SecondaryAppliesThenLosesResponse_AlsoReportsUnknownWithoutAnotherReplay()
    {
        using var provider = Build();
        provider.UpdateSite("site-1", "http://localhost:15101", "http://localhost:15102");
        var attempts = 0;
        var effects = 0;
        var error = await Assert.ThrowsAsync<global::Grpc.Core.RpcException>(() => provider.ExecuteAsync<int>("site-1", (channel, ct) =>
        {
            if (++attempts == 1) throw GrpcRetrySafety.NotDispatched("starting");
            effects++;
            throw new global::Grpc.Core.RpcException(new global::Grpc.Core.Status(global::Grpc.Core.StatusCode.Unavailable, "applied, reply lost"));
        }, CancellationToken.None));
        Assert.Equal(global::Grpc.Core.StatusCode.Aborted, error.StatusCode);
        Assert.Contains("OUTCOME_UNKNOWN", error.Status.Detail);
        Assert.Equal(2, attempts);
        Assert.Equal(1, effects);
    }

    /// <summary>Verifies that a safe read or an explicit pre-dispatch rejection can be retried on the other node.</summary>
    /// <param name="safeRead">Whether the call is a safe read rather than an explicit pre-dispatch rejection.</param>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Theory]
    [InlineData(false)]
    [InlineData(true)]
    public async Task SafeReadOrExplicitPreDispatchRejection_CanRetry(bool safeRead)
    {
        using var provider = Build();
        provider.UpdateSite("site-1", "http://localhost:15101", "http://localhost:15102");
        var attempts = 0;
        var result = await provider.ExecuteAsync("site-1", (channel, ct) =>
        {
            if (++attempts == 1)
                throw safeRead ? new global::Grpc.Core.RpcException(new global::Grpc.Core.Status(global::Grpc.Core.StatusCode.Unavailable, "lost reply"))
                    : GrpcRetrySafety.NotDispatched("not ready");
            return Task.FromResult(42);
        }, CancellationToken.None, retrySafe: safeRead);
        Assert.Equal(42, result);
        Assert.Equal(2, attempts);
    }

    /// <summary>Verifies that updating a site during an in-flight call parks the old channel and lets the call complete.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task UpdateSite_WhileCallInFlight_ParksTheOldChannelAndLetsTheCallComplete()
    {
        using var provider = Build();
        provider.UpdateSite("site-1", "http://localhost:15101", nodeBEndpoint: null);

        var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        GrpcChannel? inFlightChannel = null;

        // A long-running command (deploy) against the site's current channel.
        var call = provider.ExecuteAsync("site-1", async (channel, ct) =>
        {
            inFlightChannel = channel;
            entered.SetResult();
            await release.Task;
            // A disposed GrpcChannel throws here — this is the assertion that the address
            // refresh did not pull the channel out from under the call.
            _ = channel.CreateCallInvoker();
            return 42;
        }, CancellationToken.None);

        await entered.Task.WaitAsync(TimeSpan.FromSeconds(5));

        // The DB refresh tick re-points the site at a new address, mid-call.
        provider.UpdateSite("site-1", "http://localhost:15102", nodeBEndpoint: null);

        // The superseded channel is PARKED, not disposed synchronously.
        Assert.Equal(1, provider.ParkedChannelCount);

        release.SetResult();
        Assert.Equal(42, await call.WaitAsync(TimeSpan.FromSeconds(5)));

        // Once the grace window elapses the park drains and the channel really is disposed —
        // deferred disposal, not a leak.
        provider.SweepGraveyard(force: true);
        Assert.Equal(0, provider.ParkedChannelCount);
        Assert.Throws<ObjectDisposedException>(() => inFlightChannel!.CreateCallInvoker());
    }

    /// <summary>Verifies that removing a site parks both channels rather than disposing them inline.</summary>
    [Fact]
    public void RemoveSite_ParksBothChannelsRatherThanDisposingThemInline()
    {
        using var provider = Build();
        provider.UpdateSite("site-1", "http://localhost:15111", "http://localhost:15112");

        provider.RemoveSite("site-1");

        // Same hazard as UpdateSite: an admin deleting a site while a command is in flight.
        Assert.Equal(2, provider.ParkedChannelCount);

        provider.SweepGraveyard(force: true);
        Assert.Equal(0, provider.ParkedChannelCount);
    }

    /// <summary>Verifies that updating a site with unchanged endpoints parks no channels.</summary>
    [Fact]
    public void UpdateSite_UnchangedEndpoints_ParksNothing()
    {
        using var provider = Build();
        provider.UpdateSite("site-1", "http://localhost:15121", "http://localhost:15122");
        provider.UpdateSite("site-1", "http://localhost:15121", "http://localhost:15122");

        Assert.Equal(0, provider.ParkedChannelCount);
    }

    /// <summary>Verifies that disposing the provider drains parked channels regardless of the remaining grace period.</summary>
    [Fact]
    public void Dispose_DrainsParkedChannelsRegardlessOfRemainingGrace()
    {
        var provider = Build();
        provider.UpdateSite("site-1", "http://localhost:15131", nodeBEndpoint: null);
        provider.UpdateSite("site-1", "http://localhost:15132", nodeBEndpoint: null);
        Assert.Equal(1, provider.ParkedChannelCount);

        provider.Dispose();

        Assert.Equal(0, provider.ParkedChannelCount);
    }

    private sealed class NoKeyProvider : ISitePskProvider
    {
        /// <inheritdoc />
        public ValueTask<string> GetAsync(string siteId, CancellationToken ct) => new("k");
        /// <inheritdoc />
        public void Invalidate(string siteId) { }
    }
}