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

using Grpc.Core;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging.Abstractions;
using NSubstitute;

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

/// <summary>
/// ScadaBridge#140 — open streaming gRPC calls end on <c>ApplicationStopping</c> instead of holding
/// Kestrel's graceful drain for the whole shutdown timeout, and new streams are refused while stopping.
/// The end-to-end measurement (real Kestrel, real LocalDb sync peer, real Akka node) lives in
/// <c>IntegrationTests/GracefulStopWithOpenSyncStreamTests</c>.
/// </summary>
public sealed class ShutdownStreamCancellationInterceptorTests
{
    private static (ShutdownStreamCancellationInterceptor Interceptor, CancellationTokenSource Stopping) Create()
    {
        var stopping = new CancellationTokenSource();
        var lifetime = Substitute.For<IHostApplicationLifetime>();
        lifetime.ApplicationStopping.Returns(stopping.Token);
        return (new ShutdownStreamCancellationInterceptor(
            lifetime, NullLogger<ShutdownStreamCancellationInterceptor>.Instance), stopping);
    }

    /// <summary>An open server stream's token fires when the host starts stopping, and the call ends.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task OpenServerStream_EndsOnApplicationStopping()
    {
        var (interceptor, stopping) = Create();
        using var callCts = new CancellationTokenSource();
        var context = new FakeServerCallContext("/svc/Stream", callCts.Token);
        var entered = new TaskCompletionSource();
        ServerCallContext? seen = null;

        var call = interceptor.ServerStreamingServerHandler<string, string>(
            "req", null!, context, async (_, _, ctx) =>
            {
                seen = ctx;
                entered.SetResult();
                await Task.Delay(Timeout.Infinite, ctx.CancellationToken).ContinueWith(_ => { });
            });

        await entered.Task.WaitAsync(TimeSpan.FromSeconds(5));
        Assert.False(call.IsCompleted);
        Assert.Equal("/svc/Stream", seen!.Method);

        stopping.Cancel();

        await call.WaitAsync(TimeSpan.FromSeconds(5));
        Assert.False(callCts.IsCancellationRequested, "the call's own token is not what ended it");
    }

    /// <summary>The handler still ends when the client goes away (the call token), as before.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task OpenDuplexStream_StillEndsOnTheCallToken()
    {
        var (interceptor, _) = Create();
        using var callCts = new CancellationTokenSource();
        var context = new FakeServerCallContext("/svc/Duplex", callCts.Token);

        var call = interceptor.DuplexStreamingServerHandler<string, string>(
            null!, null!, context,
            (_, _, ctx) => Task.Delay(Timeout.Infinite, ctx.CancellationToken).ContinueWith(_ => { }));

        callCts.Cancel();
        await call.WaitAsync(TimeSpan.FromSeconds(5));
    }

    /// <summary>A client stream returns its handler's response through the wrapper.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ClientStream_ReturnsTheHandlerResponse()
    {
        var (interceptor, _) = Create();
        var context = new FakeServerCallContext("/svc/Upload", CancellationToken.None);

        var response = await interceptor.ClientStreamingServerHandler<string, string>(
            null!, context, (_, _) => Task.FromResult("done"));

        Assert.Equal("done", response);
    }

    /// <summary>Once stopping has begun, a new stream is refused with the retryable Unavailable.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task NewStream_WhileStopping_IsRefusedUnavailable()
    {
        var (interceptor, stopping) = Create();
        stopping.Cancel();
        var invoked = false;

        var ex = await Assert.ThrowsAsync<RpcException>(() => interceptor.DuplexStreamingServerHandler<string, string>(
            null!, null!, new FakeServerCallContext("/localdb_sync.v1.LocalDbSync/Sync", CancellationToken.None),
            (_, _, _) => { invoked = true; return Task.CompletedTask; }));

        Assert.Equal(StatusCode.Unavailable, ex.StatusCode);
        Assert.False(invoked);
    }

    /// <summary>Unary calls are left alone even while stopping — finishing them is what the drain is for.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task UnaryCall_WhileStopping_PassesThrough()
    {
        var (interceptor, stopping) = Create();
        stopping.Cancel();

        var reply = await interceptor.UnaryServerHandler<string, string>(
            "req", new FakeServerCallContext("/svc/Unary", CancellationToken.None), (_, _) => Task.FromResult("ok"));

        Assert.Equal("ok", reply);
    }

    /// <summary>
    /// Program.cs registers the interceptor LAST (innermost, after the auth gates) on BOTH the
    /// central and the site <c>AddGrpc</c> pipelines.
    /// </summary>
    [Fact]
    public void Program_RegistersTheInterceptorInnermost_OnBothGrpcPipelines()
    {
        var program = Path.Combine(FindRepositoryRoot(), "src", "ZB.MOM.WW.ScadaBridge.Host", "Program.cs");
        var root = CSharpSyntaxTree.ParseText(File.ReadAllText(program)).GetRoot();

        var addGrpcLambdas = root.DescendantNodes().OfType<InvocationExpressionSyntax>()
            .Where(i => i.Expression is MemberAccessExpressionSyntax { Name.Identifier.ValueText: "AddGrpc" })
            .Select(i => i.ArgumentList.Arguments.Single().Expression)
            .OfType<LambdaExpressionSyntax>()
            .ToList();
        Assert.Equal(2, addGrpcLambdas.Count);

        foreach (var lambda in addGrpcLambdas)
        {
            var added = lambda.DescendantNodes().OfType<GenericNameSyntax>()
                .Where(n => n.Identifier.ValueText == "Add" && n.Parent is MemberAccessExpressionSyntax m
                    && m.Expression.ToString().EndsWith("Interceptors", StringComparison.Ordinal))
                .Select(n => n.TypeArgumentList.Arguments.Single().ToString())
                .ToList();
            Assert.True(added.Count >= 2, $"expected an auth gate and the shutdown interceptor, got [{string.Join(", ", added)}]");
            Assert.Equal(nameof(ShutdownStreamCancellationInterceptor), added[^1]);
        }
    }

    private static string FindRepositoryRoot()
    {
        var dir = new DirectoryInfo(AppContext.BaseDirectory);
        while (dir is not null && !File.Exists(Path.Combine(dir.FullName, "ZB.MOM.WW.ScadaBridge.slnx")))
            dir = dir.Parent;
        return dir?.FullName ?? throw new InvalidOperationException("repository root not found");
    }

    /// <summary>Minimal hand-rolled <see cref="ServerCallContext"/> (grpc-dotnet ships no test context).</summary>
    private sealed class FakeServerCallContext(string method, CancellationToken token) : ServerCallContext
    {
        /// <inheritdoc />
        protected override string MethodCore => method;
        /// <inheritdoc />
        protected override string HostCore => "localhost";
        /// <inheritdoc />
        protected override string PeerCore => "ipv4:127.0.0.1:12345";
        /// <inheritdoc />
        protected override DateTime DeadlineCore => DateTime.UtcNow.AddMinutes(1);
        /// <inheritdoc />
        protected override Metadata RequestHeadersCore { get; } = [];
        /// <inheritdoc />
        protected override CancellationToken CancellationTokenCore => token;
        /// <inheritdoc />
        protected override Metadata ResponseTrailersCore { get; } = [];
        /// <inheritdoc />
        protected override Status StatusCore { get; set; }
        /// <inheritdoc />
        protected override WriteOptions? WriteOptionsCore { get; set; }
        /// <inheritdoc />
        protected override AuthContext AuthContextCore { get; } =
            new(null, new Dictionary<string, List<AuthProperty>>());
        /// <inheritdoc />
        protected override IDictionary<object, object> UserStateCore { get; } = new Dictionary<object, object>();

        /// <inheritdoc />
        protected override ContextPropagationToken CreatePropagationTokenCore(ContextPropagationOptions? options)
            => throw new NotSupportedException();

        /// <inheritdoc />
        protected override Task WriteResponseHeadersAsyncCore(Metadata responseHeaders) => Task.CompletedTask;
    }
}