ScadaBridge/tests/ZB.MOM.WW.ScadaBridge.Host.Tests/GrpcSiteTransportFailoverTests.cs

using System.Net;
using System.Net.Sockets;
using Grpc.Core;
using Grpc.Net.Client;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Hosting;
using Microsoft.AspNetCore.TestHost;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using ZB.MOM.WW.ScadaBridge.Commons.Messages.Lifecycle;
using ZB.MOM.WW.ScadaBridge.Communication;
using ZB.MOM.WW.ScadaBridge.Communication.Grpc;

namespace ZB.MOM.WW.ScadaBridge.Host.Tests;

/// <summary>
/// T1B.3: <see cref="SitePairChannelProvider"/> failover/failback + credential/deadline proof over
/// two in-process gRPC <see cref="TestServer"/>s (NodeA preferred, NodeB standby). Exercises the
/// real client stack — PSK call credentials, per-call deadline, sticky failover, no-retry on
/// <see cref="StatusCode.DeadlineExceeded"/>, and failback to the preferred node.
/// </summary>
public sealed class GrpcSiteTransportFailoverTests : IAsyncLifetime
{
    private const string SiteId = "site-1";
    private const string SiteKey = "the-site-1-key";
    private const string EndpointA = "http://node-a";
    private const string EndpointB = "http://node-b";

    private IHost _hostA = null!;
    private IHost _hostB = null!;
    private StubSiteCommandService _stubA = null!;
    private StubSiteCommandService _stubB = null!;
    private SitePairChannelProvider _provider = null!;
    private volatile bool _preferredReachable;
    private readonly Dictionary<string, Func<HttpMessageHandler>> _extraHandlers = new(StringComparer.Ordinal);

    /// <summary>Starts the stub servers and builds the channel provider under test.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public async Task InitializeAsync()
    {
        (_hostA, _stubA) = await StartServerAsync();
        (_hostB, _stubB) = await StartServerAsync();

        var handlerA = _hostA.GetTestServer().CreateHandler();
        var handlerB = _hostB.GetTestServer().CreateHandler();

        _provider = new SitePairChannelProvider(
            new FixedPskProvider(SiteKey),
            Options.Create(new CommunicationOptions()),
            NullLogger<SitePairChannelProvider>.Instance,
            handlerFactory: endpoint => endpoint == EndpointA ? handlerA
                : endpoint == EndpointB ? handlerB
                : _extraHandlers[endpoint](),
            reachabilityProbe: (_, _) => Task.FromResult(_preferredReachable));

        _provider.UpdateSite(SiteId, EndpointA, EndpointB);
    }

    /// <summary>Stops the stub servers and disposes the channel provider after the tests.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public async Task DisposeAsync()
    {
        _provider.Dispose();
        await _hostA.StopAsync();
        await _hostB.StopAsync();
        _hostA.Dispose();
        _hostB.Dispose();
    }

    private static async Task<(IHost, StubSiteCommandService)> StartServerAsync()
    {
        var stub = new StubSiteCommandService();
        var host = await new HostBuilder()
            .ConfigureWebHost(web => web
                .UseTestServer()
                .ConfigureServices(services =>
                {
                    services.AddGrpc();
                    services.AddSingleton(stub);
                })
                .Configure(app =>
                {
                    app.UseRouting();
                    app.UseEndpoints(e => e.MapGrpcService<StubSiteCommandService>());
                }))
            .StartAsync();
        return (host, stub);
    }

    private static LifecycleRequest EnableReq() =>
        SiteCommandDtoMapper.ToLifecycleRequest(
            new EnableInstanceCommand("cmd-1", "Site1.Pump1", DateTimeOffset.UtcNow));

    private Task<LifecycleReply> CallAsync(TimeSpan? deadline = null) =>
        _provider.ExecuteAsync(
            SiteId,
            (channel, ct) =>
            {
                var client = new SiteCommandService.SiteCommandServiceClient(channel);
                return client.ExecuteLifecycleAsync(
                    EnableReq(),
                    deadline: DateTime.UtcNow + (deadline ?? TimeSpan.FromSeconds(10)),
                    cancellationToken: ct).ResponseAsync;
            },
            CancellationToken.None);

    /// <summary>Verifies the case: healthy Call, hits Preferred Node A, with Psk And Site Header And Deadline.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task HealthyCall_HitsPreferredNodeA_WithPskAndSiteHeaderAndDeadline()
    {
        var reply = await CallAsync();

        Assert.Equal(LifecycleReply.ReplyOneofCase.InstanceLifecycle, reply.ReplyCase);
        Assert.Equal(1, _stubA.LifecycleCalls);
        Assert.Equal(0, _stubB.LifecycleCalls);
        Assert.Equal($"Bearer {SiteKey}", _stubA.LastAuthHeader);
        Assert.Equal(SiteId, _stubA.LastSiteHeader);
        Assert.True(_stubA.DeadlineWasSet, "the client must set a per-call deadline");
        Assert.True(_provider.IsOnPreferredNode(SiteId));
    }

    /// <summary>Verifies the case: not Dispatched, on Node A, fails Over To Node B, and Stays Sticky.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task NotDispatched_OnNodeA_FailsOverToNodeB_AndStaysSticky()
    {
        _stubA.RejectBeforeDispatch = true;

        var reply = await CallAsync();

        // Failed over: A was tried (and threw), B answered.
        Assert.Equal(LifecycleReply.ReplyOneofCase.InstanceLifecycle, reply.ReplyCase);
        Assert.Equal(1, _stubA.LifecycleCalls);
        Assert.Equal(1, _stubB.LifecycleCalls);
        Assert.False(_provider.IsOnPreferredNode(SiteId));

        // Sticky: the next call goes straight to B without re-touching A.
        await CallAsync();
        Assert.Equal(1, _stubA.LifecycleCalls);
        Assert.Equal(2, _stubB.LifecycleCalls);
    }

    /// <summary>Verifies the case: failback, returns To Preferred Node A, once Preferred Is Reachable Again.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task Failback_ReturnsToPreferredNodeA_OncePreferredIsReachableAgain()
    {
        _stubA.RejectBeforeDispatch = true;
        await CallAsync(); // flips to B
        Assert.False(_provider.IsOnPreferredNode(SiteId));

        // NodeA recovers; the failback probe reports it reachable.
        _stubA.RejectBeforeDispatch = false;
        _preferredReachable = true;
        var back = await _provider.TryFailbackAsync(SiteId, CancellationToken.None);

        Assert.True(back);
        Assert.True(_provider.IsOnPreferredNode(SiteId));

        var beforeA = _stubA.LifecycleCalls;
        await CallAsync();
        Assert.Equal(beforeA + 1, _stubA.LifecycleCalls); // next call back on A
    }

    /// <summary>Verifies the case: deadline Exceeded, on Node A, is Not Retried On Node B.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task DeadlineExceeded_OnNodeA_IsNotRetriedOnNodeB()
    {
        // A DeadlineExceeded is ambiguous (a WriteTag/Deploy/Failover may already have executed), so
        // it must NOT fail over to B. Modelled by A returning the status
        // directly, isolating the retry-decision from TestServer's own timeout mechanics.
        _stubA.StatusToThrow = StatusCode.DeadlineExceeded;

        var ex = await Assert.ThrowsAsync<RpcException>(() => CallAsync());

        Assert.Equal(StatusCode.DeadlineExceeded, ex.StatusCode);
        Assert.Equal(1, _stubA.LifecycleCalls);
        Assert.Equal(0, _stubB.LifecycleCalls); // B was never tried
        Assert.True(_provider.IsOnPreferredNode(SiteId), "a deadline must not flip stickiness");
    }

    /// <summary>Verifies the case: ambiguous Unavailable, does Not Replay Lifecycle On Peer.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task AmbiguousUnavailable_DoesNotReplayLifecycleOnPeer()
    {
        _stubA.StatusToThrow = StatusCode.Unavailable;

        var error = await Assert.ThrowsAsync<RpcException>(() => CallAsync());

        Assert.Equal(StatusCode.Aborted, error.StatusCode);
        Assert.Contains("OUTCOME_UNKNOWN", error.Status.Detail);
        Assert.Equal(1, _stubA.LifecycleCalls);
        Assert.Equal(0, _stubB.LifecycleCalls);
    }

    /// <summary>
    /// ScadaBridge#136: NodeA is a stopped node — nothing listens on its port, so the TCP connect is
    /// refused before any request byte is written. A mutation (here a lifecycle command, the same
    /// path as an instance deploy) must fail over to NodeB within the SAME call instead of returning
    /// OUTCOME_UNKNOWN, and stay sticky on B.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ConnectionRefused_OnNodeA_MutationFailsOverToNodeB_WithinTheCall()
    {
        var refused = $"http://127.0.0.1:{ClosedLoopbackPort()}";
        _extraHandlers[refused] = () => new SocketsHttpHandler();
        _provider.UpdateSite(SiteId, refused, EndpointB);

        var reply = await CallAsync();

        Assert.Equal(LifecycleReply.ReplyOneofCase.InstanceLifecycle, reply.ReplyCase);
        Assert.Equal(1, _stubB.LifecycleCalls);
        Assert.False(_provider.IsOnPreferredNode(SiteId));
    }

    /// <summary>
    /// ScadaBridge#136, the docker-rig shape: a stopped container's hostname no longer resolves, so
    /// the connect fails at name resolution. The request was never sent; the mutation fails over.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task UnresolvableNodeA_MutationFailsOverToNodeB_WithinTheCall()
    {
        const string unresolvable = "http://scadabridge-136-stopped-node.invalid:8083";
        _extraHandlers[unresolvable] = () => new SocketsHttpHandler();
        _provider.UpdateSite(SiteId, unresolvable, EndpointB);

        var reply = await CallAsync();

        Assert.Equal(LifecycleReply.ReplyOneofCase.InstanceLifecycle, reply.ReplyCase);
        Assert.Equal(1, _stubB.LifecycleCalls);
        Assert.False(_provider.IsOnPreferredNode(SiteId));
    }

    /// <summary>
    /// ScadaBridge#136 guard: the request REACHED NodeA (its stub ran the command) and the connection
    /// then died before the response — modelled by a handler that forwards to A and then throws the
    /// transport's premature-end error. That outcome is genuinely unknown: it must surface as
    /// OUTCOME_UNKNOWN and must NOT be replayed on NodeB.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ConnectionLostAfterDelivery_StaysOutcomeUnknown_AndIsNotReplayed()
    {
        const string dropsAfterSend = "http://node-a-drops-after-send";
        var inner = _hostA.GetTestServer().CreateHandler();
        _extraHandlers[dropsAfterSend] = () => new DropAfterDeliveryHandler(inner);
        _provider.UpdateSite(SiteId, dropsAfterSend, EndpointB);

        var error = await Assert.ThrowsAsync<RpcException>(() => CallAsync());

        Assert.Equal(StatusCode.Aborted, error.StatusCode);
        Assert.Contains("OUTCOME_UNKNOWN", error.Status.Detail);
        Assert.Equal(1, _stubA.LifecycleCalls);
        Assert.Equal(0, _stubB.LifecycleCalls);
    }

    /// <summary>A loopback port with nothing listening (bound, read, released).</summary>
    private static int ClosedLoopbackPort()
    {
        using var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        var port = ((IPEndPoint)listener.LocalEndpoint).Port;
        listener.Stop();
        return port;
    }

    /// <summary>Delivers the request to the real server, then fails as if the connection reset mid-response.</summary>
    private sealed class DropAfterDeliveryHandler(HttpMessageHandler inner) : DelegatingHandler(inner)
    {
        /// <inheritdoc />
        protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken ct)
        {
            using var delivered = await base.SendAsync(request, ct);
            await delivered.Content.ReadAsByteArrayAsync(ct);
            throw new HttpRequestException(HttpRequestError.ResponseEnded,
                "The response ended prematurely.", new IOException("connection reset",
                    new SocketException((int)SocketError.ConnectionReset)));
        }
    }

    /// <summary>Verifies the case: unknown Site, throws, site Channel Unavailable.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task UnknownSite_Throws_SiteChannelUnavailable()
    {
        await Assert.ThrowsAsync<SiteChannelUnavailableException>(
            () => _provider.ExecuteAsync<LifecycleReply>(
                "not-configured", (_, _) => Task.FromResult(new LifecycleReply()), CancellationToken.None));
    }

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

    /// <summary>Stub SiteCommandService recording what the client sent and letting a test steer faults.</summary>
    private sealed class StubSiteCommandService : SiteCommandService.SiteCommandServiceBase
    {
        private int _lifecycleCalls;
        /// <summary>Gets the number of lifecycle calls received.</summary>
        public int LifecycleCalls => Volatile.Read(ref _lifecycleCalls);
        /// <summary>Gets the authorization header of the last call received.</summary>
        public string? LastAuthHeader { get; private set; }
        /// <summary>Gets the site header of the last call received.</summary>
        public string? LastSiteHeader { get; private set; }
        /// <summary>Gets a value indicating whether the last call carried a deadline.</summary>
        public bool DeadlineWasSet { get; private set; }
        public volatile bool RejectBeforeDispatch;
        public StatusCode? StatusToThrow;

        /// <inheritdoc />
        public override Task<LifecycleReply> ExecuteLifecycle(LifecycleRequest request, ServerCallContext context)
        {
            Interlocked.Increment(ref _lifecycleCalls);
            LastAuthHeader = context.RequestHeaders
                .FirstOrDefault(h => h.Key == ControlPlaneCredentials.AuthorizationHeader)?.Value;
            LastSiteHeader = context.RequestHeaders
                .FirstOrDefault(h => h.Key == ControlPlaneCredentials.SiteHeader)?.Value;
            DeadlineWasSet = context.Deadline != DateTime.MaxValue;

            if (RejectBeforeDispatch)
            {
                throw GrpcRetrySafety.NotDispatched("node not ready before dispatch");
            }
            if (StatusToThrow is { } status)
            {
                throw new RpcException(new Status(status, "modelled fault"));
            }

            return Task.FromResult(SiteCommandDtoMapper.ToLifecycleReply(
                new InstanceLifecycleResponse("cmd-1", "Site1.Pump1", true, null, DateTimeOffset.UtcNow)));
        }
    }
}