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

using System.Collections.Concurrent;
using Akka.Actor;
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.Notification;
using ZB.MOM.WW.ScadaBridge.Communication;
using ZB.MOM.WW.ScadaBridge.Communication.Grpc;

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

/// <summary>
/// T1A.3: <see cref="GrpcCentralTransport"/> + <see cref="CentralChannelProvider"/> over a real
/// gRPC stack (two in-process <see cref="TestServer"/> central nodes, the real
/// <see cref="CentralControlGrpcService"/> and <see cref="CentralControlAuthInterceptor"/>). Proves
/// the sticky failover/failback policy, the PSK + site-header attachment, the per-call deadline,
/// and — the hard rule — no cross-node retry on <c>DeadlineExceeded</c>.
/// </summary>
/// <remarks>
/// A "down" node is modelled by a <see cref="ToggleHandler"/> that throws before reaching the
/// TestServer, so BOTH the unary call and the failback <c>Heartbeat</c> probe see it as
/// <c>Unavailable</c> — the honest shape of a refused connection, and the only class the transport
/// fails over on. Readiness is always set, so a node that is "up" answers everything.
/// </remarks>
public class GrpcCentralTransportTests : IAsyncLifetime
{
    private const string SiteA = "site-a";
    private const string SiteAKey = "site-a-preshared-key";
    private const string EndpointA = "http://central-a/";
    private const string EndpointB = "http://central-b/";

    private ActorSystem _system = null!;
    private CentralNode _nodeA = null!;
    private CentralNode _nodeB = null!;

    /// <inheritdoc />
    public async Task InitializeAsync()
    {
        _system = ActorSystem.Create("grpc-central-transport-test");
        _nodeA = await CentralNode.StartAsync(_system, "A", SiteA, SiteAKey, repliesToSubmit: true);
        _nodeB = await CentralNode.StartAsync(_system, "B", SiteA, SiteAKey, repliesToSubmit: true);
    }

    /// <inheritdoc />
    public async Task DisposeAsync()
    {
        await _nodeA.DisposeAsync();
        await _nodeB.DisposeAsync();
        await _system.Terminate();
    }

    private CentralChannelProvider NewProvider(string? pskKey = SiteAKey) => new(
        new[] { EndpointA, EndpointB },
        new FixedPskProvider(pskKey),
        SiteA,
        new CommunicationOptions(),
        NullLogger<CentralChannelProvider>.Instance,
        handlerFactory: HandlerFor,
        probeDeadline: TimeSpan.FromSeconds(2),
        backoffBase: TimeSpan.FromMilliseconds(50),
        backoffCap: TimeSpan.FromMilliseconds(200));

    /// <summary>
    /// When set, node A's handler short-circuits every call with this gRPC status
    /// instead of reaching its TestServer. See <see cref="GrpcStatusHandler"/>.
    /// </summary>
    private Grpc.Core.StatusCode? _nodeAForcedStatus;

    private HttpMessageHandler HandlerFor(string endpoint) => endpoint == EndpointA
        ? new GrpcStatusHandler(
            new ToggleHandler(_nodeA.Server.CreateHandler(), () => _nodeA.IsUp),
            () => _nodeAForcedStatus)
        : new ToggleHandler(_nodeB.Server.CreateHandler(), () => _nodeB.IsUp);

    private GrpcCentralTransport NewTransport(CentralChannelProvider provider, CommunicationOptions? options = null)
        => new(provider, options ?? new CommunicationOptions(), NullLogger<GrpcCentralTransport>.Instance);

    [Fact]
    public async Task HappyPath_ReachesThePreferredNode_AndRoutesTheAckBack()
    {
        using var provider = NewProvider();
        var transport = NewTransport(provider);
        var inbox = new Capture(_system);

        transport.SubmitNotification(NewSubmit("n1"), inbox.Ref);

        var ack = Assert.IsType<NotificationSubmitAck>(inbox.Receive(TimeSpan.FromSeconds(5)));
        Assert.True(ack.Accepted);
        Assert.Equal("n1", ack.NotificationId);
        Assert.Equal(0, provider.CurrentIndex); // stayed on preferred
        Assert.Equal(1, _nodeA.SubmitCount);
        Assert.Equal(0, _nodeB.SubmitCount);
    }

    [Fact]
    public async Task Sticky_StaysOnThePreferredNode_WhileHealthy()
    {
        using var provider = NewProvider();
        var transport = NewTransport(provider);

        for (var i = 0; i < 4; i++)
        {
            var inbox = new Capture(_system);
            transport.SubmitNotification(NewSubmit($"n{i}"), inbox.Ref);
            Assert.IsType<NotificationSubmitAck>(inbox.Receive(TimeSpan.FromSeconds(5)));
        }

        Assert.Equal(0, provider.CurrentIndex);
        Assert.Equal(4, _nodeA.SubmitCount);
        Assert.Equal(0, _nodeB.SubmitCount);
    }

    [Fact]
    public async Task Failover_FlipsToThePeer_WhenThePreferredIsUnavailable()
    {
        using var provider = NewProvider();
        var transport = NewTransport(provider);
        _nodeA.IsUp = false; // preferred refuses connections

        var inbox = new Capture(_system);
        transport.SubmitNotification(NewSubmit("n1"), inbox.Ref);

        var ack = Assert.IsType<NotificationSubmitAck>(inbox.Receive(TimeSpan.FromSeconds(5)));
        Assert.True(ack.Accepted);
        Assert.Equal(1, provider.CurrentIndex); // flipped to the peer
        Assert.Equal(0, _nodeA.SubmitCount);
        Assert.Equal(1, _nodeB.SubmitCount);
    }

    [Fact]
    public async Task Failback_ReturnsToThePreferred_OnceItIsReachableAgain()
    {
        using var provider = NewProvider();
        var transport = NewTransport(provider);

        // Take the preferred down and drive one call so we flip to the peer + arm the failback probe.
        _nodeA.IsUp = false;
        var inbox = new Capture(_system);
        transport.SubmitNotification(NewSubmit("n1"), inbox.Ref);
        Assert.IsType<NotificationSubmitAck>(inbox.Receive(TimeSpan.FromSeconds(5)));
        Assert.Equal(1, provider.CurrentIndex);

        // Bring the preferred back; the background probe should fail us back within a few backoffs.
        _nodeA.IsUp = true;
        await WaitUntil(() => provider.CurrentIndex == 0, TimeSpan.FromSeconds(5));
        Assert.Equal(0, provider.CurrentIndex);

        // New calls resume on the preferred node.
        var inbox2 = new Capture(_system);
        transport.SubmitNotification(NewSubmit("n2"), inbox2.Ref);
        Assert.IsType<NotificationSubmitAck>(inbox2.Receive(TimeSpan.FromSeconds(5)));
        Assert.True(_nodeA.SubmitCount >= 1);
    }

    [Fact]
    public async Task PskAndSiteHeader_AreAttached_SoTheGatedCallReachesTheService()
    {
        // The service is gated by CentralControlAuthInterceptor; a call that reaches it (and gets
        // Accepted) proves both the bearer PSK and the x-scadabridge-site header were attached.
        using var provider = NewProvider(pskKey: SiteAKey);
        var transport = NewTransport(provider);
        var inbox = new Capture(_system);

        transport.SubmitNotification(NewSubmit("n1"), inbox.Ref);

        var ack = Assert.IsType<NotificationSubmitAck>(inbox.Receive(TimeSpan.FromSeconds(5)));
        Assert.True(ack.Accepted);
    }

    [Fact]
    public async Task WrongPsk_IsRejected_AndNotRetriedOnThePeer()
    {
        // PermissionDenied is not a connect failure — the transport surfaces it as a transient
        // Status.Failure without flipping to the peer.
        using var provider = NewProvider(pskKey: "the-wrong-key");
        var transport = NewTransport(provider);
        var inbox = new Capture(_system);

        transport.SubmitNotification(NewSubmit("n1"), inbox.Ref);

        Assert.IsType<Status.Failure>(inbox.Receive(TimeSpan.FromSeconds(5)));
        Assert.Equal(0, provider.CurrentIndex);   // no flip
        Assert.Equal(0, _nodeB.SubmitCount);      // peer never tried
    }

    [Fact]
    public async Task DeadlineExceeded_IsNotRetriedOnThePeer()
    {
        // THE hard rule: on DeadlineExceeded the call may ALREADY have executed, so the transport
        // must surface Status.Failure and must NOT try node B.
        //
        // The status is injected rather than produced by black-holing node A behind a short
        // deadline. That older setup was load-dependent and failed intermittently in full-solution
        // sweeps: on a saturated machine the call could fail to even START, which IsConnectFailure
        // correctly classifies as provably-unsent, so the transport failed over and node B's ack
        // arrived instead of a Status.Failure. The test then read as a flake while actually
        // reporting that its own premise had not held. Injecting the status makes the failure mode
        // the test's subject rather than a race — see BlackHoledNode_DoesNotHang for the
        // deadline-is-actually-applied half.
        _nodeAForcedStatus = Grpc.Core.StatusCode.DeadlineExceeded;

        using var provider = NewProvider();
        var transport = NewTransport(provider);
        var inbox = new Capture(_system);

        transport.SubmitNotification(NewSubmit("n1"), inbox.Ref);

        Assert.IsType<Status.Failure>(inbox.Receive(TimeSpan.FromSeconds(5)));
        Assert.Equal(0, provider.CurrentIndex);   // no failover on a deadline
        Assert.Equal(0, _nodeB.SubmitCount);      // peer never tried
    }

    [Fact]
    public async Task BlackHoledNode_DoesNotHang_APerCallDeadlineIsApplied()
    {
        // The other half of the split: a node that accepts the call and never replies must not
        // hang the caller forever — a per-call deadline bounds it. Both nodes black-hole, so this
        // holds whichever node the transport ends up on and the assertion cannot be perturbed by
        // whether the machine was loaded enough to turn the stall into a connect failure.
        _nodeA.SetBlackHole();
        _nodeB.SetBlackHole();
        var shortDeadline = new CommunicationOptions
        {
            NotificationForwardTimeout = TimeSpan.FromMilliseconds(300),
        };

        using var provider = NewProvider();
        var transport = NewTransport(provider, shortDeadline);
        var inbox = new Capture(_system);

        transport.SubmitNotification(NewSubmit("n1"), inbox.Ref);

        // Returns rather than hanging: the 5 s inbox wait is far longer than the 300 ms deadline,
        // so a missing deadline shows up as a TimeoutException from Receive.
        Assert.IsType<Status.Failure>(inbox.Receive(TimeSpan.FromSeconds(5)));
    }

    private static NotificationSubmit NewSubmit(string id) => new(
        NotificationId: id,
        ListName: "ops",
        Subject: "s",
        Body: "b",
        SourceSiteId: SiteA,
        SourceInstanceId: null,
        SourceScript: null,
        SiteEnqueuedAt: DateTimeOffset.UtcNow);

    /// <summary>
    /// Spins until <paramref name="condition"/> holds, then asserts it — a bare
    /// <c>return</c> on timeout would make every future caller's wait silently
    /// vacuous. Today's single caller happens to re-assert immediately after,
    /// so this only sharpens the failure message; it is here so the next caller
    /// does not have to remember to.
    /// </summary>
    private static async Task WaitUntil(Func<bool> condition, TimeSpan timeout)
    {
        var deadline = DateTime.UtcNow + timeout;
        while (DateTime.UtcNow < deadline)
        {
            if (condition())
            {
                return;
            }

            await Task.Delay(25);
        }

        Assert.True(condition(), $"Condition was still false after waiting {timeout}.");
    }

    /// <summary>
    /// A raw message sink used as the transport's <c>replyTo</c>. Unlike Akka's <c>Inbox</c>, which
    /// rethrows a <see cref="Status.Failure"/>'s cause on receive, this captures every message
    /// verbatim so a test can assert on the <see cref="Status.Failure"/> itself.
    /// </summary>
    private sealed class Capture
    {
        private readonly BlockingCollection<object> _messages = new();

        public Capture(ActorSystem system)
        {
            Ref = system.ActorOf(Props.Create(() => new CaptureActor(_messages)));
        }

        public IActorRef Ref { get; }

        public object Receive(TimeSpan timeout)
            => _messages.TryTake(out var message, timeout)
                ? message
                : throw new TimeoutException("No message captured within the timeout.");

        private sealed class CaptureActor : ReceiveActor
        {
            public CaptureActor(BlockingCollection<object> messages) => ReceiveAny(messages.Add);
        }
    }

    /// <summary>
    /// Short-circuits a call with a chosen gRPC status, so a test can pick the exact failure class
    /// the transport must classify instead of trying to provoke it with timing.
    /// </summary>
    /// <remarks>
    /// Emits a trailers-only response: HTTP 200 with <c>grpc-status</c> in the HEADERS and an empty
    /// body, which is the shape gRPC defines for a call that fails before producing a message and
    /// which <c>Grpc.Net.Client</c> surfaces as an <c>RpcException</c> carrying that status.
    /// </remarks>
    private sealed class GrpcStatusHandler : DelegatingHandler
    {
        private readonly Func<Grpc.Core.StatusCode?> _forced;

        public GrpcStatusHandler(HttpMessageHandler inner, Func<Grpc.Core.StatusCode?> forced)
        {
            InnerHandler = inner;
            _forced = forced;
        }

        protected override async Task<HttpResponseMessage> SendAsync(
            HttpRequestMessage request, CancellationToken cancellationToken)
        {
            var forced = _forced();
            if (forced == null)
            {
                return await base.SendAsync(request, cancellationToken).ConfigureAwait(false);
            }

            var response = new HttpResponseMessage(System.Net.HttpStatusCode.OK)
            {
                Version = new Version(2, 0),
                Content = new ByteArrayContent(Array.Empty<byte>()),
                RequestMessage = request,
            };
            response.Content.Headers.ContentType =
                new System.Net.Http.Headers.MediaTypeHeaderValue("application/grpc");
            response.Headers.TryAddWithoutValidation(
                "grpc-status", ((int)forced.Value).ToString(System.Globalization.CultureInfo.InvariantCulture));
            response.Headers.TryAddWithoutValidation("grpc-message", "injected by GrpcStatusHandler");
            return response;
        }
    }

    /// <summary>A gRPC channel handler that throws (a refused connection) while its node is "down".</summary>
    private sealed class ToggleHandler : DelegatingHandler
    {
        private readonly Func<bool> _isUp;

        public ToggleHandler(HttpMessageHandler inner, Func<bool> isUp)
        {
            InnerHandler = inner;
            _isUp = isUp;
        }

        protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
        {
            if (!_isUp())
            {
                throw new HttpRequestException("simulated central node down");
            }

            return await base.SendAsync(request, cancellationToken).ConfigureAwait(false);
        }
    }

    private sealed class FixedPskProvider(string? key) : ISitePskProvider
    {
        public ValueTask<string> GetAsync(string siteId, CancellationToken ct)
            => key is null ? throw new InvalidOperationException("no key") : new ValueTask<string>(key);

        public void Invalidate(string siteId) { }
    }

    /// <summary>One in-process central node: TestServer + real service/interceptor + a stub actor.</summary>
    private sealed class CentralNode : IAsyncDisposable
    {
        private IHost _host = null!;
        private IActorRef _stub = null!;
        private readonly StubCounters _counters = new();

        public TestServer Server { get; private set; } = null!;
        public volatile bool IsUp = true;
        public int SubmitCount => _counters.Submits;

        public static async Task<CentralNode> StartAsync(
            ActorSystem system, string label, string site, string key, bool repliesToSubmit)
        {
            var node = new CentralNode();
            node._stub = system.ActorOf(
                Props.Create(() => new StubCentralActor(node._counters, repliesToSubmit)), $"stub-{label}");

            var service = new CentralControlGrpcService(
                NullLogger<CentralControlGrpcService>.Instance,
                Options.Create(new CommunicationOptions()));
            service.SetReady(node._stub);

            var psk = new MapPskProvider(new Dictionary<string, string> { [site] = key });

            node._host = await new HostBuilder()
                .ConfigureWebHost(web => web
                    .UseTestServer()
                    .ConfigureServices(services =>
                    {
                        services.AddGrpc(o => o.Interceptors.Add<CentralControlAuthInterceptor>());
                        services.AddSingleton<ISitePskProvider>(psk);
                        services.AddSingleton(service);
                    })
                    .Configure(app =>
                    {
                        app.UseRouting();
                        app.UseEndpoints(e => e.MapGrpcService<CentralControlGrpcService>());
                    }))
                .StartAsync();

            node.Server = node._host.GetTestServer();
            return node;
        }

        /// <summary>Switches the node's actor to a black hole that counts but never replies.</summary>
        public void SetBlackHole() => _counters.BlackHole = true;

        public async ValueTask DisposeAsync()
        {
            await _host.StopAsync();
            _host.Dispose();
        }

        private sealed class StubCounters
        {
            private int _submits;
            public int Submits => Volatile.Read(ref _submits);
            public void IncrementSubmits() => Interlocked.Increment(ref _submits);
            public volatile bool BlackHole;
        }

        private sealed class StubCentralActor : ReceiveActor
        {
            public StubCentralActor(StubCounters counters, bool repliesToSubmit)
            {
                Receive<NotificationSubmit>(msg =>
                {
                    counters.IncrementSubmits();
                    if (repliesToSubmit && !counters.BlackHole)
                    {
                        Sender.Tell(new NotificationSubmitAck(msg.NotificationId, Accepted: true, Error: null));
                    }
                });

                // Heartbeat lands here as a Tell (the failback probe); ignore it, no reply expected.
                ReceiveAny(_ => { });
            }
        }

        private sealed class MapPskProvider(IReadOnlyDictionary<string, string> keys) : ISitePskProvider
        {
            public ValueTask<string> GetAsync(string siteId, CancellationToken ct)
                => keys.TryGetValue(siteId, out var key)
                    ? new ValueTask<string>(key)
                    : throw new InvalidOperationException($"no key for '{siteId}'");

            public void Invalidate(string siteId) { }
        }
    }
}