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

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;

    /// <inheritdoc />
    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 : handlerB,
            reachabilityProbe: (_, _) => Task.FromResult(_preferredReachable));

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

    /// <inheritdoc />
    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);

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

    [Fact]
    public async Task Unavailable_OnNodeA_FailsOverToNodeB_AndStaysSticky()
    {
        _stubA.ThrowUnavailable = 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);
    }

    [Fact]
    public async Task Failback_ReturnsToPreferredNodeA_OncePreferredIsReachableAgain()
    {
        _stubA.ThrowUnavailable = true;
        await CallAsync(); // flips to B
        Assert.False(_provider.IsOnPreferredNode(SiteId));

        // NodeA recovers; the failback probe reports it reachable.
        _stubA.ThrowUnavailable = 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
    }

    [Fact]
    public async Task DeadlineExceeded_OnNodeA_IsNotRetriedOnNodeB()
    {
        // A DeadlineExceeded is ambiguous (a WriteTag/Deploy/Failover may already have executed), so
        // — unlike Unavailable — 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");
    }

    [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
    {
        public ValueTask<string> GetAsync(string siteId, CancellationToken ct) => new(key);
        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;
        public int LifecycleCalls => Volatile.Read(ref _lifecycleCalls);
        public string? LastAuthHeader { get; private set; }
        public string? LastSiteHeader { get; private set; }
        public bool DeadlineWasSet { get; private set; }
        public volatile bool ThrowUnavailable;
        public StatusCode? StatusToThrow;

        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 (ThrowUnavailable)
            {
                throw new RpcException(new Status(StatusCode.Unavailable, "node down"));
            }
            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)));
        }
    }
}