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

using Shouldly;
using Xunit;
using ZB.MOM.WW.OtOpcUa.Core.Abstractions;
using ZB.MOM.WW.OtOpcUa.Driver.Galaxy.Browse;
using ZB.MOM.WW.OtOpcUa.Driver.Galaxy.Health;
using ZB.MOM.WW.OtOpcUa.Driver.Galaxy.Runtime;
using static ZB.MOM.WW.OtOpcUa.Driver.Galaxy.Tests.Runtime.GalaxyDriverFixRound1Tests;

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

/// <summary>
///     Lane F fix round 2 (re-review of #735):
///     <list type="bullet">
///       <item>N1 — the init platform seed runs in the background: a slow or cancelled repository
///       fetch can neither delay nor fail <c>InitializeAsync</c> (Galaxy's Tier-A Initialize budget is
///       5 s, and a cold gateway cache can spend all of it), and a seed that lands after a teardown
///       is discarded.</item>
///       <item>N2 — a seed requested under a new runtime always gets a pass of its own, even when a
///       worker from the previous runtime is still running.</item>
///       <item>N3 — a subscribe / subscribe-once read whose SubscribeBulk straddles a teardown fails
///       and unsubscribes instead of registering into the next runtime's registry.</item>
///       <item>The probe watcher's own subscribe re-checks its stop after the await.</item>
///     </list>
/// </summary>
public sealed class GalaxyDriverFixRound2Tests
{
    private static readonly TimeSpan InitBudget = TimeSpan.FromSeconds(2);

    /// <summary>A platform source whose fetches park until released (optionally ignoring cancellation).</summary>
    private sealed class GatedPlatformSource : IGalaxyPlatformSource
    {
        private int _calls;
        private readonly TaskCompletionSource _release = new(TaskCreationOptions.RunContinuationsAsynchronously);

        public int Calls => Volatile.Read(ref _calls);
        public int ParkCall { get; init; } = 1;
        public bool IgnoreCancellation { get; init; }
        public Func<int, IReadOnlyList<string>> Result { get; init; } = _ => ["Plat01"];
        public TaskCompletionSource Entered { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
        public TaskCompletionSource<CancellationToken> ParkedToken { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
        public void Release() => _release.TrySetResult();

        public async Task<IReadOnlyList<string>> GetPlatformTagNamesAsync(CancellationToken cancellationToken)
        {
            var call = Interlocked.Increment(ref _calls);
            if (call == ParkCall)
            {
                ParkedToken.TrySetResult(cancellationToken);
                Entered.TrySetResult();
                if (IgnoreCancellation) await _release.Task;
                else await _release.Task.WaitAsync(cancellationToken);
            }
            return Result(call);
        }
    }

    private sealed class ThrowingPlatformSource(Exception ex) : IGalaxyPlatformSource
    {
        private int _calls;
        public int Calls => Volatile.Read(ref _calls);

        public Task<IReadOnlyList<string>> GetPlatformTagNamesAsync(CancellationToken cancellationToken)
        {
            Interlocked.Increment(ref _calls);
            return Task.FromException<IReadOnlyList<string>>(ex);
        }
    }

    private static async Task Entered(TaskCompletionSource tcs, string because) =>
        (await Task.WhenAny(tcs.Task, Task.Delay(WaitMs))).ShouldBe(tcs.Task, because);

    // ───────────────────────────── N1: background init seed ─────────────────────────────

    /// <summary>
    ///     A repository fetch that never completes cannot hold InitializeAsync: init returns Healthy
    ///     inside its budget, and the teardown cancels the parked fetch.
    /// </summary>
    [Fact]
    public async Task InitializeAsync_returns_Healthy_while_the_platform_fetch_never_completes()
    {
        var sub = new ParkableSubscriber();
        var platforms = new GatedPlatformSource();
        await using var driver = new GalaxyDriver("drv", Opts(), hierarchySource: null, subscriber: sub, platformSource: platforms);
        try
        {
            var sw = System.Diagnostics.Stopwatch.StartNew();
            await driver.InitializeAsync("{}", CancellationToken.None).WaitAsync(InitBudget);
            sw.Stop();

            driver.GetHealth().State.ShouldBe(DriverState.Healthy);
            await Entered(platforms.Entered, "positive control: the seed fetch was started");
            sw.Elapsed.ShouldBeLessThan(InitBudget);

            await driver.ShutdownAsync(CancellationToken.None).WaitAsync(TimeSpan.FromMilliseconds(WaitMs));
            var token = await platforms.ParkedToken.Task;
            token.IsCancellationRequested.ShouldBeTrue("the teardown must cancel a seed still in flight");
        }
        finally
        {
            platforms.Release();
        }
    }

    /// <summary>A fetch that is cancelled (the gateway's own deadline) does not fail init.</summary>
    [Fact]
    public async Task A_cancelled_platform_fetch_does_not_fail_InitializeAsync()
    {
        var sub = new ParkableSubscriber();
        var platforms = new ThrowingPlatformSource(new OperationCanceledException("gateway deadline"));
        await using var driver = new GalaxyDriver("drv", Opts(), hierarchySource: null, subscriber: sub, platformSource: platforms);

        await driver.InitializeAsync("{}", CancellationToken.None).WaitAsync(InitBudget);

        driver.GetHealth().State.ShouldBe(DriverState.Healthy);
        await WaitUntilAsync(() => platforms.Calls == 1, "positive control: the seed fetch ran");
        driver.ProbeWatcherForTests.ShouldNotBeNull().WatchedPlatforms.ShouldBeEmpty();

        // The driver keeps serving after the failed seed.
        await driver.SubscribeAsync(["Tank01.Level"], TimeSpan.Zero, CancellationToken.None);
        driver.TrackedSubscriptionCountForTests.ShouldBe(1);
    }

    /// <summary>
    ///     The production failure mode: the host's Tier-A Initialize timeout cancels the init token
    ///     while the (cold-cache) fetch is still running. Init must already have returned Healthy, and
    ///     the seed — tied to the runtime, not to the init call — must survive that cancellation and
    ///     still land.
    /// </summary>
    [Fact]
    public async Task Init_token_cancellation_mid_fetch_neither_fails_init_nor_cancels_the_seed()
    {
        var sub = new ParkableSubscriber();
        var platforms = new GatedPlatformSource();
        await using var driver = new GalaxyDriver("drv", Opts(), hierarchySource: null, subscriber: sub, platformSource: platforms);
        using var initCts = new CancellationTokenSource();
        try
        {
            initCts.CancelAfter(TimeSpan.FromMilliseconds(300)); // the Polly Initialize timeout, shrunk
            await driver.InitializeAsync("{}", initCts.Token).WaitAsync(InitBudget);
            driver.GetHealth().State.ShouldBe(DriverState.Healthy);

            await Entered(platforms.Entered, "the seed fetch must be in flight");
            await WaitUntilAsync(() => initCts.IsCancellationRequested, "the init token fires while the fetch is parked");
            (await platforms.ParkedToken.Task).IsCancellationRequested.ShouldBeFalse(
                "the seed belongs to the runtime; the init call's timeout must not cancel it");
        }
        finally
        {
            platforms.Release();
        }

        await WaitUntilAsync(() => driver.ProbeWatcherForTests?.WatchedPlatforms.Count == 1,
            "the seed must still land after the init token was cancelled");
    }

    /// <summary>Positive control: a fetch that completes after init still populates host statuses.</summary>
    [Fact]
    public async Task A_fetch_completing_after_init_populates_host_statuses()
    {
        var sub = new ParkableSubscriber();
        var platforms = new GatedPlatformSource();
        await using var driver = new GalaxyDriver("drv", Opts(), hierarchySource: null, subscriber: sub, platformSource: platforms);
        try
        {
            await driver.InitializeAsync("{}", CancellationToken.None).WaitAsync(InitBudget);
            await Entered(platforms.Entered, "the seed must be parked");
            driver.ProbeWatcherForTests.ShouldNotBeNull().WatchedPlatforms.ShouldBeEmpty();
        }
        finally
        {
            platforms.Release();
        }

        await WaitUntilAsync(() => driver.ProbeWatcherForTests?.WatchedPlatforms.Count == 1,
            "the late seed must reach the probe watcher");
        await driver.SubscribeAsync(["Tank01.Level"], TimeSpan.Zero, CancellationToken.None); // starts the pump
        await sub.EmitBoolAsync(sub.HandleByRef["Plat01.ScanState"], true);
        await WaitUntilAsync(() => StateOf(driver, "Plat01") == HostState.Running,
            "a ScanState event must surface the late-seeded platform");
    }

    // ───────────────────────────── N1 generation + N2: straddling worker ─────────────────────────────

    /// <summary>
    ///     A seed from the previous runtime that completes after a teardown + re-init is discarded,
    ///     and the new runtime's own seed request still gets its own pass.
    /// </summary>
    [Fact]
    public async Task A_seed_straddling_a_teardown_is_discarded_and_the_new_runtime_is_seeded()
    {
        var sub = new ParkableSubscriber();
        var platforms = new GatedPlatformSource
        {
            IgnoreCancellation = true, // a fetch that finishes regardless — only the generation check can stop it
            Result = call => call == 1 ? ["Stale01"] : ["Plat01"],
        };
        await using var driver = new GalaxyDriver("drv", Opts(), hierarchySource: null, subscriber: sub, platformSource: platforms);
        try
        {
            await driver.InitializeAsync("{}", CancellationToken.None).WaitAsync(InitBudget);
            await Entered(platforms.Entered, "runtime 1's seed must be parked");

            await driver.ShutdownAsync(CancellationToken.None).WaitAsync(TimeSpan.FromMilliseconds(WaitMs));
            await driver.InitializeAsync("{}", CancellationToken.None).WaitAsync(InitBudget);
        }
        finally
        {
            platforms.Release();
        }

        await WaitUntilAsync(() => platforms.Calls == 2, "the request made under the new runtime must get its own pass");
        await WaitUntilAsync(() => driver.ProbeWatcherForTests?.WatchedPlatforms.Count == 1,
            "the new runtime must be seeded");
        driver.ProbeWatcherForTests!.WatchedPlatforms.ShouldBe(["Plat01"]);
        lock (sub.Subscribed) sub.Subscribed.ShouldNotContain("Stale01.ScanState", "runtime 1's late result must be discarded");
    }

    // ───────────────────────────── probe watcher: re-check after the SubscribeBulk await ─────────────────────────────

    /// <summary>
    ///     A watcher sync whose SubscribeBulk completes after StopAsync must not write the probe entry
    ///     back into the registry, and must release the handles it just acquired. Positive control:
    ///     the same sync without a stop registers.
    /// </summary>
    [Fact]
    public async Task Watcher_sync_straddling_a_stop_does_not_resurrect_the_probe_entry()
    {
        {
            var registry = new SubscriptionRegistry();
            using var watcher = new PerPlatformProbeWatcher(new ParkableSubscriber(), new HostStatusAggregator(), registry);
            await watcher.SyncPlatformsAsync(["A"], CancellationToken.None);
            registry.Contains(SubscriptionRegistry.ProbeSubscriptionId).ShouldBeTrue("positive control");
        }

        {
            var registry = new SubscriptionRegistry();
            var sub = new ParkableSubscriber { ParkSubscribeCall = 1 };
            var watcher = new PerPlatformProbeWatcher(sub, new HostStatusAggregator(), registry);

            var sync = watcher.SyncPlatformsAsync(["A"], CancellationToken.None);
            await Entered(sub.SubscribeEntered, "the sync must be parked in SubscribeBulk");
            await watcher.StopAsync(CancellationToken.None);

            sub.ReleaseSubscribe();
            await sync.WaitAsync(TimeSpan.FromMilliseconds(WaitMs));

            registry.Contains(SubscriptionRegistry.ProbeSubscriptionId).ShouldBeFalse("a stopped watcher must not re-register");
            watcher.WatchedPlatforms.ShouldBeEmpty();
            lock (sub.Unsubscribed) sub.Unsubscribed.ShouldContain(sub.HandleByRef["A.ScanState"],
                "the handle acquired after the stop must be released");
        }
    }

    // ───────────────────────────── N3: subscribe / read straddling a teardown ─────────────────────────────

    /// <summary>
    ///     A SubscribeAsync whose SubscribeBulk completes after the teardown fails, registers nothing
    ///     and unsubscribes its handles; a subscribe on the live runtime registers (positive control).
    /// </summary>
    [Fact]
    public async Task SubscribeAsync_straddling_a_teardown_fails_and_registers_nothing()
    {
        var sub = new ParkableSubscriber { ParkSubscribeCall = 2 };
        await using var driver = new GalaxyDriver("drv", Opts(), hierarchySource: null, subscriber: sub);
        await driver.InitializeAsync("{}", CancellationToken.None);

        await driver.SubscribeAsync(["Tank01.Level"], TimeSpan.Zero, CancellationToken.None);
        driver.TrackedSubscriptionCountForTests.ShouldBe(1, "positive control: a live-runtime subscribe registers");

        var straddling = driver.SubscribeAsync(["Tank02.Level"], TimeSpan.Zero, CancellationToken.None);
        try
        {
            await Entered(sub.SubscribeEntered, "the subscribe must be parked in SubscribeBulk");
            await driver.ShutdownAsync(CancellationToken.None).WaitAsync(TimeSpan.FromMilliseconds(WaitMs));
            driver.TrackedSubscriptionCountForTests.ShouldBe(0);
        }
        finally
        {
            sub.ReleaseSubscribe();
        }

        await Should.ThrowAsync<InvalidOperationException>(() => straddling.WaitAsync(TimeSpan.FromMilliseconds(WaitMs)));
        driver.TrackedSubscriptionCountForTests.ShouldBe(0, "the straddling subscribe must not land in the registry");
        lock (sub.Unsubscribed) sub.Unsubscribed.ShouldContain(sub.HandleByRef["Tank02.Level"]);

        await driver.InitializeAsync("{}", CancellationToken.None);
        driver.TrackedSubscriptionCountForTests.ShouldBe(0, "the next runtime starts with an empty registry");
    }

    /// <summary>The subscribe-once read path: same straddle, same outcome.</summary>
    [Fact]
    public async Task Subscribe_once_read_straddling_a_teardown_fails_and_registers_nothing()
    {
        var sub = new ParkableSubscriber { ParkSubscribeCall = 1 };
        await using var driver = new GalaxyDriver("drv", Opts(), hierarchySource: null, subscriber: sub);
        await driver.InitializeAsync("{}", CancellationToken.None);

        var read = driver.ReadAsync(["Tank01.Level"], CancellationToken.None);
        try
        {
            await Entered(sub.SubscribeEntered, "the read must be parked in SubscribeBulk");
            await driver.ShutdownAsync(CancellationToken.None).WaitAsync(TimeSpan.FromMilliseconds(WaitMs));
        }
        finally
        {
            sub.ReleaseSubscribe();
        }

        await Should.ThrowAsync<InvalidOperationException>(() => read.WaitAsync(TimeSpan.FromMilliseconds(WaitMs)));
        driver.TrackedSubscriptionCountForTests.ShouldBe(0);
        lock (sub.Unsubscribed) sub.Unsubscribed.ShouldContain(sub.HandleByRef["Tank01.Level"]);
    }
}