ScadaBridge/tests/ZB.MOM.WW.ScadaBridge.ManagementService.Tests/DebugStreamHubSubscribeRaceTests.cs

using Microsoft.AspNetCore.Http;
using Microsoft.AspNetCore.Http.Connections.Features;
using Microsoft.AspNetCore.Http.Features;
using Microsoft.AspNetCore.SignalR;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using NSubstitute;
using ZB.MOM.WW.ScadaBridge.Commons.Entities.Instances;
using ZB.MOM.WW.ScadaBridge.Commons.Interfaces.Repositories;
using ZB.MOM.WW.ScadaBridge.Commons.Messages.DebugView;
using ZB.MOM.WW.ScadaBridge.Commons.Messages.Streaming;
using ZB.MOM.WW.ScadaBridge.Communication;
using ZB.MOM.WW.ScadaBridge.Communication.Grpc;

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

/// <summary>
/// Regression tests for the hub-side subscribe race. <c>SubscribeInstance</c>
/// awaits <c>StartStreamAsync</c> — up to 30s — and only then stores the session id on the
/// connection. A client that disconnects inside that window ran <c>OnDisconnectedAsync</c>
/// while there was still no session id to release, so the session stayed registered in the
/// singleton <see cref="DebugStreamService"/> and its keepalive renewed it forever (the
/// orphan reaper only collects UNregistered sessions). The fix publishes the id and
/// re-checks connection liveness under the per-connection <c>Items</c> monitor, so exactly
/// one side tears the session down.
/// </summary>
public class DebugStreamHubSubscribeRaceTests
{
    private sealed class GatedDebugStreamService : DebugStreamService
    {
        private readonly TaskCompletionSource _gate = new(TaskCreationOptions.RunContinuationsAsynchronously);
        private readonly TaskCompletionSource _entered = new(TaskCreationOptions.RunContinuationsAsynchronously);
        private readonly List<string> _stopped = new();

        /// <summary>Creates the gated debug stream service with an unreleased start gate.</summary>
        public GatedDebugStreamService()
            : base(
                new CommunicationService(Options.Create(new CommunicationOptions()), NullLogger<CommunicationService>.Instance),
                new ServiceCollection().BuildServiceProvider(),
                new SiteStreamGrpcClientFactory(NullLoggerFactory.Instance),
                NullLogger<DebugStreamService>.Instance)
        {
        }

        /// <summary>The fixed session id returned by every started stream.</summary>
        public string SessionId => "session-under-test";

        /// <summary>Completes once <see cref="StartStreamAsync"/> has been entered.</summary>
        public Task Entered => _entered.Task;

        /// <summary>Session ids passed to <see cref="StopStream"/>, in call order.</summary>
        public IReadOnlyList<string> Stopped
        {
            get { lock (_stopped) { return _stopped.ToList(); } }
        }

        /// <summary>Releases the start gate so a pending <c>StartStreamAsync</c> call completes.</summary>
        public void ReleaseStart() => _gate.TrySetResult();

        /// <inheritdoc />
        public override async Task<DebugStreamSession> StartStreamAsync(
            int instanceId, Action<object> onEvent, Action onTerminated, CancellationToken ct = default)
        {
            _entered.TrySetResult();
            await _gate.Task;
            return new DebugStreamSession(
                SessionId,
                new DebugViewSnapshot(
                    "Inst-1",
                    Array.Empty<AttributeValueChanged>(),
                    Array.Empty<AlarmStateChanged>(),
                    DateTimeOffset.UtcNow));
        }

        /// <inheritdoc />
        public override void StopStream(string sessionId)
        {
            lock (_stopped) { _stopped.Add(sessionId); }
        }
    }

    private sealed class TestHubCallerContext : HubCallerContext
    {
        private readonly FeatureCollection _features = new();
        private readonly CancellationTokenSource _aborted = new();

        /// <summary>Creates a hub caller context over the supplied HTTP context.</summary>
        /// <param name="http">The HTTP context exposed through the hub context.</param>
        public TestHubCallerContext(HttpContext http) =>
            _features.Set<IHttpContextFeature>(new TestHttpContextFeature { HttpContext = http });

        /// <inheritdoc />
        public override string ConnectionId => "test-conn";
        /// <inheritdoc />
        public override string? UserIdentifier => null;
        /// <inheritdoc />
        public override System.Security.Claims.ClaimsPrincipal? User => null;
        /// <inheritdoc />
        public override IDictionary<object, object?> Items { get; } = new Dictionary<object, object?>();
        /// <inheritdoc />
        public override IFeatureCollection Features => _features;
        /// <inheritdoc />
        public override CancellationToken ConnectionAborted => _aborted.Token;
        /// <inheritdoc />
        public override void Abort() => _aborted.Cancel();

        private sealed class TestHttpContextFeature : IHttpContextFeature
        {
            /// <summary>Gets or sets the HTTP context exposed by the feature.</summary>
            public HttpContext? HttpContext { get; set; }
        }
    }

    private const string SessionIdKey = "DebugStreamSessionId";

    private static (DebugStreamHub Hub, TestHubCallerContext Context, GatedDebugStreamService Stream) CreateHub()
    {
        var repo = Substitute.For<ITemplateEngineRepository>();
        repo.GetInstanceByIdAsync(7).Returns(new Instance("Inst-1") { Id = 7, SiteId = 1 });

        var services = new ServiceCollection();
        services.AddSingleton(repo);
        var http = new DefaultHttpContext { RequestServices = services.BuildServiceProvider() };

        var stream = new GatedDebugStreamService();
        var hub = new DebugStreamHub(
            stream, Substitute.For<IHubContext<DebugStreamHub>>(), NullLogger<DebugStreamHub>.Instance)
        {
            Clients = Substitute.For<IHubCallerClients>()
        };

        var context = new TestHubCallerContext(http);
        hub.Context = context;
        // No roles / no permitted sites on the connection == system-wide scope, so the
        // authorization branch is not what these tests are exercising.
        return (hub, context, stream);
    }

    /// <summary>Verifies that subscribe instance, client disconnects mid start, stops the just created session.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task SubscribeInstance_ClientDisconnectsMidStart_StopsTheJustCreatedSession()
    {
        var (hub, context, stream) = CreateHub();

        var subscribe = hub.SubscribeInstance(7);
        await stream.Entered.WaitAsync(TimeSpan.FromSeconds(5));

        // The client goes away while the start is still in flight — the disconnect path
        // finds nothing registered yet.
        await hub.OnDisconnectedAsync(null);
        Assert.Empty(stream.Stopped);

        stream.ReleaseStart();
        await subscribe.WaitAsync(TimeSpan.FromSeconds(10));

        Assert.Equal(new[] { stream.SessionId }, stream.Stopped);
        Assert.False(context.Items.ContainsKey(SessionIdKey));
    }

    /// <summary>Verifies that subscribe instance, connection aborted mid start, stops the just created session.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task SubscribeInstance_ConnectionAbortedMidStart_StopsTheJustCreatedSession()
    {
        // Same window, but the abort token trips without OnDisconnectedAsync having run yet
        // (SignalR aborts the connection before dispatching the disconnect handler).
        var (hub, context, stream) = CreateHub();

        var subscribe = hub.SubscribeInstance(7);
        await stream.Entered.WaitAsync(TimeSpan.FromSeconds(5));

        context.Abort();
        stream.ReleaseStart();
        await subscribe.WaitAsync(TimeSpan.FromSeconds(10));

        Assert.Equal(new[] { stream.SessionId }, stream.Stopped);
        Assert.False(context.Items.ContainsKey(SessionIdKey));
    }

    /// <summary>Verifies that subscribe instance, then disconnect, stops the session exactly once.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task SubscribeInstance_ThenDisconnect_StopsTheSessionExactlyOnce()
    {
        var (hub, context, stream) = CreateHub();

        stream.ReleaseStart();
        await hub.SubscribeInstance(7);

        Assert.Empty(stream.Stopped);
        Assert.Equal(stream.SessionId, context.Items[SessionIdKey]);

        await hub.OnDisconnectedAsync(null);
        Assert.Equal(new[] { stream.SessionId }, stream.Stopped);
        Assert.False(context.Items.ContainsKey(SessionIdKey));

        // A second release (an explicit unsubscribe racing the disconnect) is a no-op.
        await hub.UnsubscribeInstance();
        Assert.Equal(new[] { stream.SessionId }, stream.Stopped);
    }
}