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

using Akka.Actor;
using Akka.TestKit.Xunit2;
using Grpc.Core;
using Microsoft.Extensions.Logging.Abstractions;
using NSubstitute;
using ZB.MOM.WW.ScadaBridge.Communication.Grpc;

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

/// <summary>
/// Verifies that after a gRPC stream is cancelled, the SiteStreamManager
/// subscription is properly cleaned up with no leaked subscriptions.
/// </summary>
// Serialized with SiteStreamGrpcServerTests (register row 102) — this class's real
// SiteStreamGrpcServer streams move the process-global subscribers.active gauge that
// SiteStreamSubscribersActiveGauge_GoesToOneOnConnect_AndBackToZeroOnCancel baselines.
[Collection("SiteStreamSubscriberGauge")]
public class CleanupVerificationTests : TestKit
{
    /// <summary>Verifies that cancelling a stream cleans up its site stream manager subscription.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task Stream_Cancellation_CleansUp_SiteStreamManager_Subscription()
    {
        // Arrange: create server with mock subscriber that tracks subscribe/remove calls
        var subscriber = Substitute.For<ISiteStreamSubscriber>();
        var subscribeCalled = false;
        var removeCalled = false;
        IActorRef? subscribedActor = null;

        subscriber.Subscribe(Arg.Any<string>(), Arg.Any<IActorRef>())
            .Returns(ci =>
            {
                subscribeCalled = true;
                subscribedActor = ci.Arg<IActorRef>();
                return "sub-cleanup-test";
            });

        subscriber.When(x => x.RemoveSubscriber(Arg.Any<IActorRef>()))
            .Do(_ => removeCalled = true);

        var logger = NullLogger<SiteStreamGrpcServer>.Instance;
        var server = new SiteStreamGrpcServer(subscriber, logger);
        server.SetReady(Sys);

        var cts = new CancellationTokenSource();
        var context = Substitute.For<ServerCallContext>();
        context.CancellationToken.Returns(cts.Token);

        var writer = Substitute.For<IServerStreamWriter<SiteStreamEvent>>();

        var request = new InstanceStreamRequest
        {
            CorrelationId = "corr-cleanup-verify",
            InstanceUniqueName = "Site1.TestInst"
        };

        // Act: start a stream, wait for it to register, then cancel
        var streamTask = Task.Run(() => server.SubscribeInstance(request, writer, context));

        await WaitForConditionAsync(() => subscribeCalled);
        Assert.True(subscribeCalled, "Subscribe should have been called");
        Assert.Equal(1, server.ActiveStreamCount);

        cts.Cancel();
        await streamTask;

        // Assert: verify cleanup
        Assert.True(removeCalled, "RemoveSubscriber should have been called after cancellation");
        Assert.Equal(0, server.ActiveStreamCount);

        // Verify the same actor that was subscribed is the one that was removed
        subscriber.Received(1).RemoveSubscriber(subscribedActor!);
    }

    /// <summary>Verifies that cancelling several streams cleans up all of their subscriptions.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task Multiple_Streams_Cancelled_AllCleanedUp()
    {
        var subscriber = Substitute.For<ISiteStreamSubscriber>();
        var removeCount = 0;

        subscriber.Subscribe(Arg.Any<string>(), Arg.Any<IActorRef>())
            .Returns("sub-multi");

        subscriber.When(x => x.RemoveSubscriber(Arg.Any<IActorRef>()))
            .Do(_ => Interlocked.Increment(ref removeCount));

        var logger = NullLogger<SiteStreamGrpcServer>.Instance;
        var server = new SiteStreamGrpcServer(subscriber, logger);
        server.SetReady(Sys);

        // Start 3 streams
        var ctsList = new List<CancellationTokenSource>();
        var tasks = new List<Task>();

        for (var i = 0; i < 3; i++)
        {
            var cts = new CancellationTokenSource();
            ctsList.Add(cts);
            var ctx = Substitute.For<ServerCallContext>();
            ctx.CancellationToken.Returns(cts.Token);
            var w = Substitute.For<IServerStreamWriter<SiteStreamEvent>>();

            var req = new InstanceStreamRequest
            {
                CorrelationId = $"corr-multi-{i}",
                InstanceUniqueName = $"Site1.Inst{i}"
            };

            tasks.Add(Task.Run(() => server.SubscribeInstance(req, w, ctx)));
        }

        await WaitForConditionAsync(() => server.ActiveStreamCount == 3);

        // Cancel all
        foreach (var cts in ctsList)
            cts.Cancel();

        await Task.WhenAll(tasks);

        Assert.Equal(0, server.ActiveStreamCount);
        Assert.Equal(3, removeCount);
    }

    private static async Task WaitForConditionAsync(Func<bool> condition, int timeoutMs = 5000)
    {
        var deadline = DateTime.UtcNow.AddMilliseconds(timeoutMs);
        while (!condition() && DateTime.UtcNow < deadline)
        {
            await Task.Delay(25);
        }

        Assert.True(condition(), $"Condition not met within {timeoutMs}ms");
    }
}