lmxopcua/tests/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Galaxy.Tests/Runtime/StreamAttachedAckTests.cs

using System.Runtime.CompilerServices;
using System.Threading.Channels;
using Google.Protobuf.WellKnownTypes;
using Microsoft.Extensions.Logging;
using Shouldly;
using Xunit;
using ZB.MOM.WW.MxGateway.Contracts.Proto;
using ZB.MOM.WW.OtOpcUa.Core.Abstractions;
using ZB.MOM.WW.OtOpcUa.Driver.Galaxy.Config;
using ZB.MOM.WW.OtOpcUa.Driver.Galaxy.Runtime;

namespace ZB.MOM.WW.OtOpcUa.Driver.Galaxy.Tests.Runtime;

/// <summary>
///     #877 — the gateway's stream-attached acknowledgement (mxaccessgw#227). With the
///     <c>stream-events-attached-ack</c> capability, an advise waits until the gateway confirms the event
///     stream is REGISTERED on the session, closing the #865 attach race that the "stream requested" wait only
///     narrowed. Without it the old best-effort wait stays the fallback.
/// </summary>
/// <remarks>
///     The fake models the race the real gateway has: a stream call is issued at once but the gateway
///     registers its subscriber only later (here: when the test opens the per-stream attach gate). An advise's
///     initial value pushed before that registration is lost, exactly as on the gateway — so an advise that
///     does not wait for the ack loses its value, and one that does cannot.
/// </remarks>
public sealed class StreamAttachedAckTests
{
    private static GalaxyDriverOptions Opts() => new(
        new GalaxyGatewayOptions("https://mxgw.test:5001", "key"),
        new GalaxyMxAccessOptions("StreamAttachedAck"),
        new GalaxyRepositoryOptions(WatchDeployEvents: false),
        new GalaxyReconnectOptions());

    private static MxEvent AckSentinel(ulong liveAfter = 0) => new()
    {
        // Shaped as the gateway emits it: no family, no body, worker_sequence 0.
        Family = MxEventFamily.Unspecified,
        StreamAttached = new StreamAttached { LiveAfterSequence = liveAfter },
    };

    /// <summary>
    ///     Gateway-like subscriber whose stream registration is DELAYED until the test opens that stream's
    ///     attach gate. Records "stream:N", "attached:N" and "subscribe:&lt;refs&gt;" in order.
    /// </summary>
    private sealed class DelayedAttachSubscriber(bool supportsAck, bool sendsAck = true) : IGalaxySubscriber
    {
        private readonly Lock _gate = new();
        private readonly List<TaskCompletionSource> _attachGates = [];
        private Channel<MxEvent>? _attached;
        private int _nextHandle;
        private int _streams;

        public bool SupportsStreamAttachedAck => supportsAck;

        public int DroppedInitialValues;
        public volatile bool NextValue = true;

        public List<string> Calls { get; } = [];

        public string[] CallsSnapshot() { lock (_gate) return [.. Calls]; }

        /// <summary>Lets stream N (1-based) register on the "gateway".</summary>
        public void Attach(int stream) => GateFor(stream).TrySetResult();

        private TaskCompletionSource GateFor(int stream)
        {
            lock (_gate)
            {
                while (_attachGates.Count < stream)
                    _attachGates.Add(new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously));
                return _attachGates[stream - 1];
            }
        }

        public Task<IReadOnlyList<SubscribeResult>> SubscribeBulkAsync(
            IReadOnlyList<string> fullReferences, int bufferedUpdateIntervalMs, CancellationToken cancellationToken)
        {
            var results = new List<SubscribeResult>();
            lock (_gate)
            {
                Calls.Add("subscribe:" + string.Join(",", fullReferences));
                foreach (var r in fullReferences)
                {
                    var h = ++_nextHandle;
                    results.Add(new SubscribeResult { TagAddress = r, ItemHandle = h, WasSuccessful = true });
                    var initial = new MxEvent
                    {
                        Family = MxEventFamily.OnDataChange,
                        ItemHandle = h,
                        WorkerSequence = (ulong)h,
                        Value = new MxValue { BoolValue = NextValue },
                        Quality = 192,
                        SourceTimestamp = Timestamp.FromDateTime(DateTime.UtcNow),
                    };
                    if (_attached is null || !_attached.Writer.TryWrite(initial)) Interlocked.Increment(ref DroppedInitialValues);
                }
            }
            return Task.FromResult<IReadOnlyList<SubscribeResult>>(results);
        }

        public Task UnsubscribeBulkAsync(IReadOnlyList<int> itemHandles, CancellationToken cancellationToken) => Task.CompletedTask;

        public async IAsyncEnumerable<MxEvent> StreamEventsAsync([EnumeratorCancellation] CancellationToken cancellationToken)
        {
            var n = Interlocked.Increment(ref _streams);
            lock (_gate) Calls.Add("stream:" + n);
            // The gateway registers the subscriber only after auth etc. — asynchronously to the call.
            await GateFor(n).Task.WaitAsync(cancellationToken).ConfigureAwait(false);
            var mine = Channel.CreateUnbounded<MxEvent>();
            lock (_gate)
            {
                _attached = mine;
                Calls.Add("attached:" + n);
            }
            try
            {
                // The gateway sends the sentinel only AFTER registration, and only when asked.
                if (supportsAck && sendsAck) yield return AckSentinel();
                await foreach (var ev in mine.Reader.ReadAllAsync(cancellationToken).ConfigureAwait(false))
                {
                    yield return ev;
                }
            }
            finally
            {
                lock (_gate) { if (ReferenceEquals(_attached, mine)) _attached = null; }
            }
        }
    }

    private static string Describe(DelayedAttachSubscriber sub) =>
        $"calls=[{string.Join(" | ", sub.CallsSnapshot())}] dropped={sub.DroppedInitialValues}";

    /// <summary>
    ///     The key test. Capability present: the subscribe is NOT issued while the stream is requested but not
    ///     yet registered; it goes out only after the ack, so its initial value arrives. Fail-before: with the
    ///     ack wait skipped (pre-#877 behaviour) the subscribe goes out at "stream requested", its initial value
    ///     is pushed to no stream, and this test fails on the ordering and the dropped count.
    /// </summary>
    [Fact]
    public async Task Capability_present_subscribe_waits_for_the_ack_and_gets_its_initial_value()
    {
        var sub = new DelayedAttachSubscriber(supportsAck: true);
        await using var driver = new GalaxyDriver("drv", Opts(), hierarchySource: null, subscriber: sub)
        {
            StreamRequestWait = TimeSpan.FromSeconds(20), // generous: the gate below decides, not the budget
        };
        await driver.InitializeAsync("{}", CancellationToken.None);
        var delivered = new TaskCompletionSource<DataChangeEventArgs>(TaskCreationOptions.RunContinuationsAsynchronously);
        driver.OnDataChange += (_, e) => delivered.TrySetResult(e);

        var subscribe = driver.SubscribeAsync(["Tank01.Level"], TimeSpan.Zero, CancellationToken.None);

        await TestWait.UntilAsync(() => sub.CallsSnapshot().Contains("stream:1"),
            "positive control: the subscribe must start the event stream", () => Describe(sub));
        // The stream is requested but not registered. Give a skipped wait every chance to advise now.
        await Task.Delay(300, TestContext.Current.CancellationToken);
        sub.CallsSnapshot().ShouldNotContain(c => c.StartsWith("subscribe:", StringComparison.Ordinal),
            $"no advise may be issued before the gateway acknowledged the stream; {Describe(sub)}");
        subscribe.IsCompleted.ShouldBeFalse("the subscribe must still be waiting for the attach ack");

        sub.Attach(1);
        await TestWait.ForAsync(subscribe, "the subscribe proceeds once the ack arrives", () => Describe(sub));

        var calls = sub.CallsSnapshot();
        Array.IndexOf(calls, "attached:1").ShouldBeLessThan(
            Array.FindIndex(calls, c => c.StartsWith("subscribe:", StringComparison.Ordinal)),
            $"the advise must follow the registration; {Describe(sub)}");
        (await TestWait.ForAsync(delivered.Task, "the initial value must arrive", () => Describe(sub)))
            .FullReference.ShouldBe("Tank01.Level");
        sub.DroppedInitialValues.ShouldBe(0);
        ((EventPump)driver.EventPumpInstanceForTests!).StreamAttachedReceived.ShouldBeTrue();
    }

    /// <summary>
    ///     Capability absent (an older gateway): no ack is requested or awaited — the advise proceeds as soon as
    ///     the stream call has started (the #865 best-effort wait), even though the stream is not registered yet.
    ///     That leaves the narrowed race in place (the value is dropped here, by construction of the fake); the
    ///     point is that the driver does not hang waiting for an ack an old gateway never sends.
    /// </summary>
    [Fact]
    public async Task Capability_absent_falls_back_to_the_stream_requested_wait()
    {
        var sub = new DelayedAttachSubscriber(supportsAck: false);
        await using var driver = new GalaxyDriver("drv", Opts(), hierarchySource: null, subscriber: sub)
        {
            StreamRequestWait = TimeSpan.FromSeconds(20),
        };
        await driver.InitializeAsync("{}", CancellationToken.None);

        var sw = System.Diagnostics.Stopwatch.StartNew();
        await TestWait.ForAsync(
            driver.SubscribeAsync(["Tank01.Level"], TimeSpan.Zero, CancellationToken.None),
            "without the capability the subscribe must not wait for an ack", () => Describe(sub));
        sw.Elapsed.ShouldBeLessThan(TimeSpan.FromSeconds(10), "the 20 s budget must not have been spent on an ack");

        var pump = (EventPump)driver.EventPumpInstanceForTests!;
        pump.AttachAckExpected.ShouldBeFalse();
        pump.StreamAttached.IsCompleted.ShouldBeFalse("no ack was requested, none arrived");
        var calls = sub.CallsSnapshot();
        Array.IndexOf(calls, "stream:1").ShouldBe(0, $"the stream is still requested first; {Describe(sub)}");
        calls.ShouldContain(c => c.StartsWith("subscribe:", StringComparison.Ordinal));
        calls.ShouldNotContain("attached:1", "the subscribe went out before registration — the old, narrowed race");
        sub.Attach(1); // let the stream finish attaching so disposal is clean
    }

    /// <summary>
    ///     Capability advertised but the ack never comes (a gateway bug or a stuck attach): the advise does not
    ///     hang — it proceeds after <see cref="GalaxyDriver.StreamRequestWait"/> and logs a Warning naming the
    ///     missing acknowledgement.
    /// </summary>
    [Fact]
    public async Task Missing_ack_times_out_with_a_warning_instead_of_hanging()
    {
        var sub = new DelayedAttachSubscriber(supportsAck: true, sendsAck: false);
        var log = new CapturingLogger();
        await using var driver = new GalaxyDriver("drv", Opts(), hierarchySource: null, subscriber: sub,
            logger: new TypedLogger(log))
        {
            StreamRequestWait = TimeSpan.FromMilliseconds(300),
        };
        await driver.InitializeAsync("{}", CancellationToken.None);
        sub.Attach(1); // registered, but no sentinel is sent

        await TestWait.ForAsync(
            driver.SubscribeAsync(["Tank01.Level"], TimeSpan.Zero, CancellationToken.None),
            "a missing ack must not hang the subscribe", () => Describe(sub));

        lock (log.Entries)
        {
            log.Entries.ShouldContain(e => e.Level == LogLevel.Warning && e.Text.Contains("attach acknowledgement", StringComparison.Ordinal),
                $"a missing ack must be logged loudly; log=\n{log.AllText}");
        }
        sub.CallsSnapshot().ShouldContain(c => c.StartsWith("subscribe:", StringComparison.Ordinal));
    }

    /// <summary>
    ///     The sentinel is a control signal: never dispatched, never counted as a received event, and it
    ///     completes <see cref="EventPump.StreamAttached"/>. Positive control: a real event after it IS dispatched.
    /// </summary>
    [Fact]
    public async Task Ack_sentinel_is_never_dispatched()
    {
        var subscriber = new ChannelSubscriber();
        var registry = new SubscriptionRegistry();
        registry.Register(1, [new TagBinding("Tank.Level", ItemHandle: 7)]);

        var dispatched = 0;
        var realValue = new TaskCompletionSource<DataChangeEventArgs>(TaskCreationOptions.RunContinuationsAsynchronously);
        await using var pump = new EventPump(subscriber, registry, channelCapacity: 8, clientName: "AckTest");
        pump.OnDataChange += (_, args) => { Interlocked.Increment(ref dispatched); realValue.TrySetResult(args); };
        pump.Start();

        // An ack shaped with OnDataChange + item handle 7 would be dispatched if it reached the channel —
        // so this proves interception, not just "unknown family is dropped".
        var ack = AckSentinel(liveAfter: 0);
        ack.Family = MxEventFamily.OnDataChange;
        ack.ItemHandle = 7;
        ack.Value = new MxValue { DoubleValue = -1 };
        await subscriber.EmitAsync(ack);
        await TestWait.ForAsync(pump.StreamAttached, "the ack must complete the attached signal");
        pump.StreamAttachedReceived.ShouldBeTrue();

        await subscriber.EmitAsync(new MxEvent
        {
            Family = MxEventFamily.OnDataChange,
            ItemHandle = 7,
            Value = new MxValue { DoubleValue = 12.5 },
            Quality = 192,
            SourceTimestamp = Timestamp.FromDateTime(DateTime.UtcNow),
        });

        (await TestWait.ForAsync(realValue.Task, "positive control: real events still dispatch")).Snapshot.Value.ShouldBe(12.5);
        Volatile.Read(ref dispatched).ShouldBe(1, "only the real event may be dispatched, never the ack");
    }

    /// <summary>
    ///     A pump whose loop ends before any ack completes <see cref="EventPump.StreamAttached"/> (so no waiter
    ///     hangs) WITHOUT reporting the ack as received.
    /// </summary>
    [Fact]
    public async Task A_dead_stream_releases_ack_waiters_without_claiming_an_ack()
    {
        await using var pump = new EventPump(
            new ThrowingSubscriber(), new SubscriptionRegistry(), onStreamFault: _ => { });
        pump.Start();
        await TestWait.ForAsync(pump.StreamAttached, "a dead pump must not strand an ack waiter");
        pump.HasEnded.ShouldBeTrue();
        pump.StreamAttachedReceived.ShouldBeFalse();
        pump.AttachAckExpected.ShouldBeTrue();
    }

    /// <summary>
    ///     Reconnect: the replay's new stream gets a FRESH attach signal — the first stream's ack does not
    ///     release it. The replay's re-advise waits for the second stream's ack, and its new initial value
    ///     arrives. Positive control: before stream 2 attaches, the re-advise has not been issued.
    /// </summary>
    [Fact]
    public async Task Reconnect_gives_each_stream_a_fresh_signal()
    {
        var sub = new DelayedAttachSubscriber(supportsAck: true);
        await using var driver = new GalaxyDriver("drv", Opts(), hierarchySource: null, subscriber: sub)
        {
            StreamRequestWait = TimeSpan.FromSeconds(20),
        };
        await driver.InitializeAsync("{}", CancellationToken.None);
        var last = new System.Collections.Concurrent.ConcurrentDictionary<string, object?>();
        driver.OnDataChange += (_, e) => last[e.FullReference] = e.Snapshot.Value;

        sub.Attach(1);
        await driver.SubscribeAsync(["Tank01.Level"], TimeSpan.Zero, CancellationToken.None);
        await TestWait.UntilAsync(() => last.TryGetValue("Tank01.Level", out var v) && v is true,
            "positive control: the first stream delivers the initial value", () => Describe(sub));
        var pump1 = (EventPump)driver.EventPumpInstanceForTests!;
        pump1.StreamAttached.IsCompleted.ShouldBeTrue();

        sub.NextValue = false;
        var replay = driver.InvokeReplayForTestAsync(CancellationToken.None);
        await TestWait.UntilAsync(() => sub.CallsSnapshot().Contains("stream:2"),
            "the replay must start a new stream", () => Describe(sub));
        var pump2 = (EventPump)driver.EventPumpInstanceForTests!;
        pump2.ShouldNotBeSameAs(pump1);
        pump2.StreamAttached.IsCompleted.ShouldBeFalse("a new stream starts with an uncompleted attach signal");

        await Task.Delay(300, TestContext.Current.CancellationToken);
        sub.CallsSnapshot().Count(c => c.StartsWith("subscribe:", StringComparison.Ordinal))
            .ShouldBe(1, $"the re-advise must wait for stream 2's ack; {Describe(sub)}");

        sub.Attach(2);
        await TestWait.ForAsync(replay, "the replay completes once stream 2 is acknowledged", () => Describe(sub));
        await TestWait.UntilAsync(() => last.TryGetValue("Tank01.Level", out var v) && v is false,
            "the re-advised tag's new initial value must arrive on stream 2", () => Describe(sub));
        sub.DroppedInitialValues.ShouldBe(0);
    }

    /// <summary>A stream fed directly by the test; reports the ack capability.</summary>
    private sealed class ChannelSubscriber : IGalaxySubscriber
    {
        private readonly Channel<MxEvent> _stream = Channel.CreateUnbounded<MxEvent>();

        public bool SupportsStreamAttachedAck => true;

        public Task<IReadOnlyList<SubscribeResult>> SubscribeBulkAsync(
            IReadOnlyList<string> fullReferences, int bufferedUpdateIntervalMs, CancellationToken cancellationToken)
            => Task.FromResult<IReadOnlyList<SubscribeResult>>([]);

        public Task UnsubscribeBulkAsync(IReadOnlyList<int> itemHandles, CancellationToken cancellationToken) => Task.CompletedTask;

        public IAsyncEnumerable<MxEvent> StreamEventsAsync(CancellationToken cancellationToken)
            => _stream.Reader.ReadAllAsync(cancellationToken);

        public ValueTask EmitAsync(MxEvent ev) => _stream.Writer.WriteAsync(ev);
    }

    private sealed class ThrowingSubscriber : IGalaxySubscriber
    {
        public bool SupportsStreamAttachedAck => true;

        public Task<IReadOnlyList<SubscribeResult>> SubscribeBulkAsync(
            IReadOnlyList<string> fullReferences, int bufferedUpdateIntervalMs, CancellationToken cancellationToken)
            => Task.FromResult<IReadOnlyList<SubscribeResult>>([]);

        public Task UnsubscribeBulkAsync(IReadOnlyList<int> itemHandles, CancellationToken cancellationToken) => Task.CompletedTask;

        public IAsyncEnumerable<MxEvent> StreamEventsAsync(CancellationToken cancellationToken)
            => throw new InvalidOperationException("not connected");
    }

    private sealed class TypedLogger(CapturingLogger inner) : ILogger<GalaxyDriver>
    {
        public IDisposable? BeginScope<TState>(TState state) where TState : notnull => null;

        public bool IsEnabled(LogLevel logLevel) => true;

        public void Log<TState>(LogLevel logLevel, EventId eventId, TState state, Exception? exception,
            Func<TState, Exception?, string> formatter)
        {
            lock (inner.Entries) inner.Log(logLevel, eventId, state, exception, formatter);
        }
    }
}