lmxopcua/tests/Drivers/ZB.MOM.WW.OtOpcUa.Driver.MTConnect.Tests/MTConnectPumpGenerationTests.cs

using Shouldly;
using Xunit;
using ZB.MOM.WW.OtOpcUa.Core.Abstractions;

namespace ZB.MOM.WW.OtOpcUa.Driver.MTConnect.Tests;

/// <summary>
///     STAB-20 / #602 — a Subscribe that lands between <c>StopSampleStreamAsync</c> clearing the pump slots
///     and the teardown that retires the client started a pump on the soon-to-be-disposed session. The
///     lifecycle restart then saw "a pump is running" and never started one against the new session, leaving
///     the subscription with no live stream. The pump is now stamped with the session generation and a
///     start request supersedes a pump from an older one.
/// </summary>
public sealed class MTConnectPumpGenerationTests
{
    private static readonly TimeSpan Watchdog = TimeSpan.FromSeconds(10);

    private static CancellationToken Ct => TestContext.Current.CancellationToken;

    /// <summary>The pump a mid-restart Subscribe started on the retired client is replaced by one on the new client.</summary>
    [Fact]
    public async Task Subscribe_during_a_lifecycle_restart_rebinds_the_pump_to_the_new_client()
    {
        var clientA = CannedAgentClient.FromFixtures();
        var clientB = CannedAgentClient.FromFixtures();
        var clients = new Queue<CannedAgentClient>([clientA, clientB]);

        var driver = new MTConnectDriver(
            new MTConnectDriverOptions
            {
                AgentUri = "http://fixture-agent:5000",
                RequestTimeoutMs = 5000,
                Tags = [new MTConnectTagDefinition("dev1_pos", DriverDataType.Float64)],
            },
            "mt-gen",
            _ => clients.Dequeue());
        await driver.InitializeAsync("{}", Ct);
        Task[] retired = [];

        var blockOnce = 0;
        var blocking = new ManualResetEventSlim(false);
        var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        driver.OnDataChange += (_, e) =>
        {
            // Once only: the second subscriber's initial-data callback replays the same 1.5 and must not wedge.
            if (e.Snapshot.Value is not 1.5d || Interlocked.Exchange(ref blockOnce, 1) != 0)
            {
                return;
            }

            entered.TrySetResult();
            blocking.Wait();
        };

        try
        {
            await driver.SubscribeAsync(["dev1_pos"], TimeSpan.FromMilliseconds(50), Ct);
            _ = clientA.PumpAsync(MTConnectPumpGenerationChunks.Pos(108, 113, "1.5"));
            await entered.Task.WaitAsync(Watchdog, Ct);

            // The first pump is now wedged in caller code. A re-initialize clears the pump slots and then
            // waits (bounded) on it, which is the window the race lives in.
            var restart = driver.InitializeAsync("{}", Ct);
            var spins = 0;
            while (driver.SampleStreamTask is not null)
            {
                await Task.Delay(10, Ct);
                (++spins).ShouldBeLessThan(500, "restart never cleared the pump slot");
            }

            // A second subscriber arrives mid-restart: it starts a pump on the (still-live) old client.
            await driver.SubscribeAsync(["dev1_pos"], TimeSpan.FromMilliseconds(50), Ct);
            await clientA.WaitForSampleCallsAsync(2).WaitAsync(Watchdog, Ct);
            restart.IsCompleted.ShouldBeFalse();

            blocking.Set();
            await restart.WaitAsync(Watchdog, Ct);

            // The pump retired by the restart is retained so shutdown can join it (it used to be dropped
            // un-awaited and could write health after the new session was Healthy).
            retired = driver.SupersededPumpTasks;
            retired.ShouldNotBeEmpty();

            // The restart installed a new client; the standing subscription must be streaming from IT.
            await clientB.WaitForSampleCallsAsync(1).WaitAsync(Watchdog, Ct);
            clientA.IsDisposed.ShouldBeTrue();
        }
        finally
        {
            blocking.Set();
            await driver.ShutdownAsync(CancellationToken.None);
        }

        // Shutdown joined every retired pump and forgot them.
        retired.ShouldAllBe(t => t.IsCompleted);
        driver.SupersededPumpTasks.ShouldBeEmpty();
    }
}

internal static class MTConnectPumpGenerationChunks
{
    public static MTConnectStreamsResult Pos(long first, long next, string value) =>
        new(
            1655000000L,
            next,
            first,
            [new MTConnectObservation("dev1_pos", value, new DateTime(2026, 7, 24, 12, 30, 0, DateTimeKind.Utc))]);
}