lmxopcua/tests/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Mqtt.IntegrationTests/MqttPrimaryHostLiveTests.cs

using System.Buffers;
using System.Collections.Concurrent;
using MQTTnet;
using Shouldly;
using Xunit;
using ZB.MOM.WW.OtOpcUa.Core.Abstractions;
using ZB.MOM.WW.OtOpcUa.Driver.Mqtt.Sparkplug;

namespace ZB.MOM.WW.OtOpcUa.Driver.Mqtt.IntegrationTests;

/// <summary>
///     Live gate for the Sparkplug primary-host claim (Gitea #511), against the Mosquitto fixture.
/// </summary>
/// <remarks>
///     <para>
///         The issue's acceptance criteria, leg by leg: the broker holds a retained <c>ONLINE</c> while
///         the claim runs; an abrupt death leaves <c>OFFLINE</c> (the Will fires); a clean shutdown
///         leaves <c>OFFLINE</c>; and the pair-node question is answered rather than left to chance —
///         a node that is not Primary claims nothing, and a role change moves the claim.
///     </para>
///     <para>
///         <b>Each test uses its own host id.</b> A Sparkplug host id is a broker-wide singleton by
///         definition, and the retained STATE outlives the test that wrote it — two tests sharing an id
///         would see each other's leftovers, and would do so intermittently depending on ordering.
///     </para>
///     <para>
///         <b>Retained messages are cleaned up.</b> A retained <c>ONLINE</c> for a host that no longer
///         exists is exactly the lie this feature prevents, so leaving one on a shared fixture would be
///         poor manners at best and a confusing false positive for the next run at worst.
///     </para>
/// </remarks>
[Trait("Category", "Integration")]
[Collection("MqttFixture")]
public sealed class MqttPrimaryHostLiveTests(MqttFixture fixture)
{
    private static readonly TimeSpan Timeout = TimeSpan.FromSeconds(15);

    private void SkipIfUnconfigured() => Assert.SkipWhen(fixture.SkipReason is not null, fixture.SkipReason ?? "");

    [Fact]
    public async Task A_primary_publishes_a_retained_ONLINE_that_a_late_subscriber_still_sees()
    {
        // Retention is the point, not an optimisation: an edge node that connects AFTER the host came
        // up must still learn the host is alive. A non-retained ONLINE reaches only whoever happened to
        // be subscribed at that instant, and every later joiner sees an absent host forever.
        SkipIfUnconfigured();

        var hostId = NewHostId();
        var ct = TestContext.Current.CancellationToken;
        await using var host = NewHost(hostId, isPrimary: true);

        await host.StartAsync(ct);
        host.IsClaimed.ShouldBeTrue("the fixture broker accepted the claim connection");

        // Subscribe only NOW — after the birth was published — so only a retained message can satisfy it.
        await using var watcher = await StateWatcher.StartAsync(fixture, hostId, ct);
        var body = await watcher.WaitForOneAsync(Timeout, ct);

        SparkplugStatePayload.TryDecode(body, out var online, out var timestamp).ShouldBeTrue();
        online.ShouldBeTrue();
        timestamp.ShouldBeGreaterThan(0);

        await host.StopAsync(ct);
        await ClearRetainedAsync(hostId, ct);
    }

    [Fact]
    public async Task An_abrupt_death_leaves_a_retained_OFFLINE_because_the_will_fires()
    {
        // THE leg the whole design exists for, and the one no unit test can reach: proof that the
        // broker is holding our death certificate and delivers it when the socket dies without a
        // goodbye. Without it a crashed node would leave a retained ONLINE asserting forever that a
        // dead host is alive — strictly worse than never publishing at all.
        SkipIfUnconfigured();

        var hostId = NewHostId();
        var ct = TestContext.Current.CancellationToken;
        var host = NewHost(hostId, isPrimary: true);

        await host.StartAsync(ct);
        host.IsClaimed.ShouldBeTrue();

        await using var watcher = await StateWatcher.StartAsync(fixture, hostId, ct);
        var birth = await watcher.WaitForOneAsync(Timeout, ct);
        SparkplugStatePayload.TryDecode(birth, out var wasOnline, out var birthTimestamp).ShouldBeTrue();
        wasOnline.ShouldBeTrue();

        // Kill the socket with no DISCONNECT. A clean teardown would make the broker DISCARD the Will
        // by protocol rule, so this test would then be observing the explicit stand-down publish and
        // would pass against an implementation that never armed a Will.
        await host.AbandonClaimForTests();

        var death = await watcher.WaitForMatchAsync(
            p => SparkplugStatePayload.TryDecode(p, out var o, out _) && !o, Timeout, ct);

        SparkplugStatePayload.TryDecode(death, out _, out var deathTimestamp).ShouldBeTrue();

        // Sparkplug 3.0 pairs a host's birth and death by timestamp — an edge node may ignore an
        // OFFLINE older than the ONLINE it holds. Mismatched timestamps produce two well-formed
        // messages and a death certificate conformant hosts are entitled to throw away.
        deathTimestamp.ShouldBe(birthTimestamp);

        await host.DisposeAsync();
        await ClearRetainedAsync(hostId, ct);
    }

    [Fact]
    public async Task A_clean_shutdown_leaves_a_retained_OFFLINE_without_waiting_for_the_will()
    {
        SkipIfUnconfigured();

        var hostId = NewHostId();
        var ct = TestContext.Current.CancellationToken;
        var host = NewHost(hostId, isPrimary: true);

        await host.StartAsync(ct);
        await using var watcher = await StateWatcher.StartAsync(fixture, hostId, ct);
        await watcher.WaitForMatchAsync(
            p => SparkplugStatePayload.TryDecode(p, out var o, out _) && o, Timeout, ct);

        await host.StopAsync(ct);

        // Promptly — the assertion is that this arrives well inside the keep-alive window it would take
        // the broker to notice a dead socket, which is why the explicit publish exists at all.
        await watcher.WaitForMatchAsync(
            p => SparkplugStatePayload.TryDecode(p, out var o, out _) && !o, Timeout, ct);

        await host.DisposeAsync();
        await ClearRetainedAsync(hostId, ct);
    }

    [Fact]
    public async Task A_node_that_is_not_primary_publishes_nothing_at_all()
    {
        // The pair-node answer. Not merely "does not win the race" — a Secondary must never open a
        // claim connection, because the Will it would register is the hazard: its own death would then
        // publish OFFLINE over a perfectly healthy Primary's ONLINE and stop the whole group.
        SkipIfUnconfigured();

        var hostId = NewHostId();
        var ct = TestContext.Current.CancellationToken;
        await using var watcher = await StateWatcher.StartAsync(fixture, hostId, ct);
        await using var host = NewHost(hostId, isPrimary: false);

        await host.StartAsync(ct);
        host.ClaimAttemptsForTests.ShouldBe(0);

        // Give the reconcile loop several ticks to misbehave in.
        await Task.Delay(TimeSpan.FromSeconds(3), ct);

        host.IsClaimed.ShouldBeFalse();
        host.ClaimAttemptsForTests.ShouldBe(0);
        watcher.Count.ShouldBe(0, "a Secondary must not touch the STATE topic");
    }

    [Fact]
    public async Task A_failover_moves_the_claim_offline_then_online_again()
    {
        // Both halves of the accepted trade, observed: losing Primary publishes OFFLINE (a genuine
        // state transition — for a window there really is no host), and regaining it republishes
        // ONLINE. The republish is the half that would be easy to omit; without it a node that failed
        // back would run as Primary while every edge node believed it was gone.
        SkipIfUnconfigured();

        var hostId = NewHostId();
        var ct = TestContext.Current.CancellationToken;
        var role = new MutableRole(isPrimary: true);
        await using var host = new SparkplugPrimaryHost(Options(hostId), "live-failover", role);

        await using var watcher = await StateWatcher.StartAsync(fixture, hostId, ct);

        await host.StartAsync(ct);
        await watcher.WaitForMatchAsync(
            p => SparkplugStatePayload.TryDecode(p, out var o, out _) && o, Timeout, ct);

        role.IsPrimary = false;
        await watcher.WaitForMatchAsync(
            p => SparkplugStatePayload.TryDecode(p, out var o, out _) && !o, Timeout, ct);
        await WaitUntilAsync(() => !host.IsClaimed, ct);

        var offlineSeen = watcher.CountMatching(
            p => SparkplugStatePayload.TryDecode(p, out var o, out _) && !o);

        role.IsPrimary = true;
        await WaitUntilAsync(() => host.IsClaimed, ct);

        // A SECOND online, after the offline — not merely "an online exists", which the birth already
        // satisfied before the failover even started.
        await WaitUntilAsync(
            () => watcher.CountMatching(p => SparkplugStatePayload.TryDecode(p, out var o, out _) && o) >= 2,
            ct);

        offlineSeen.ShouldBeGreaterThanOrEqualTo(1);

        await host.StopAsync(ct);
        await ClearRetainedAsync(hostId, ct);
    }

    // ---- helpers -----------------------------------------------------------------------------------

    /// <summary>A per-test host id: the retained STATE outlives the test that wrote it.</summary>
    private static string NewHostId() => $"otopcua-live-{Guid.NewGuid():N}"[..28];

    private MqttDriverOptions Options(string hostId) => new()
    {
        Host = fixture.PlainHost,
        Port = fixture.PlainPort,
        Username = fixture.Username,
        Password = fixture.Password,
        UseTls = false,
        ClientId = $"otopcua-live-{Guid.NewGuid():N}"[..20],
        ConnectTimeoutSeconds = 10,
        Mode = MqttMode.SparkplugB,
        Sparkplug = new MqttSparkplugOptions
        {
            GroupId = "OtOpcUaSim",
            HostId = hostId,
            ActAsPrimaryHost = true,
        },
    };

    private SparkplugPrimaryHost NewHost(string hostId, bool isPrimary) =>
        new(Options(hostId), $"live-{hostId}", new MutableRole(isPrimary));

    /// <summary>
    ///     Clears the retained STATE this test left behind — a zero-length retained publish is MQTT's
    ///     "forget this". Leaving a retained ONLINE for a host that no longer exists is precisely the
    ///     lie the feature prevents.
    /// </summary>
    private async Task ClearRetainedAsync(string hostId, CancellationToken ct)
    {
        var client = new MqttClientFactory().CreateMqttClient();
        try
        {
            await client.ConnectAsync(
                new MqttClientOptionsBuilder()
                    .WithTcpServer(fixture.PlainHost, fixture.PlainPort)
                    .WithCredentials(fixture.Username, fixture.Password)
                    .WithClientId($"otopcua-live-clean-{Guid.NewGuid():N}"[..23])
                    .WithTimeout(TimeSpan.FromSeconds(10))
                    .Build(),
                ct);

            await client.PublishAsync(
                new MqttApplicationMessageBuilder()
                    .WithTopic(SparkplugTopic.FormatState(hostId))
                    .WithPayload(Array.Empty<byte>())
                    .WithRetainFlag(true)
                    .Build(),
                ct);

            await client.DisconnectAsync(cancellationToken: ct);
        }
        finally
        {
            client.Dispose();
        }
    }

    private static async Task WaitUntilAsync(Func<bool> condition, CancellationToken ct)
    {
        var deadline = DateTime.UtcNow + Timeout;
        while (!condition())
        {
            if (DateTime.UtcNow >= deadline)
            {
                throw new TimeoutException($"Condition was not satisfied within {Timeout}.");
            }

            await Task.Delay(100, ct);
        }
    }

    /// <summary>A role source a test can flip, standing in for a redundancy failover.</summary>
    private sealed class MutableRole(bool isPrimary) : IDriverPrimaryRoleSource
    {
        private volatile bool _isPrimary = isPrimary;

        public bool IsPrimary
        {
            get => _isPrimary;
            set => _isPrimary = value;
        }

        public bool ShouldActAsPrimary => _isPrimary;

        public void PublishPrimaryRole(bool shouldActAsPrimary) => _isPrimary = shouldActAsPrimary;
    }

    /// <summary>Subscribes one host's STATE topic and records every body seen, in order.</summary>
    private sealed class StateWatcher : IAsyncDisposable
    {
        private readonly IMqttClient _client;
        private readonly ConcurrentQueue<byte[]> _bodies = new();

        private StateWatcher(IMqttClient client) => _client = client;

        public int Count => _bodies.Count;

        public static async Task<StateWatcher> StartAsync(MqttFixture fixture, string hostId, CancellationToken ct)
        {
            var client = new MqttClientFactory().CreateMqttClient();
            var watcher = new StateWatcher(client);

            client.ApplicationMessageReceivedAsync += args =>
            {
                watcher._bodies.Enqueue(args.ApplicationMessage.Payload.ToArray());
                return Task.CompletedTask;
            };

            await client.ConnectAsync(
                new MqttClientOptionsBuilder()
                    .WithTcpServer(fixture.PlainHost, fixture.PlainPort)
                    .WithCredentials(fixture.Username, fixture.Password)
                    .WithClientId($"otopcua-live-state-{Guid.NewGuid():N}"[..23])
                    .WithTimeout(TimeSpan.FromSeconds(10))
                    .Build(),
                ct);

            await client.SubscribeAsync(
                new MqttClientSubscribeOptionsBuilder()
                    .WithTopicFilter(f => f.WithTopic(SparkplugTopic.FormatState(hostId)))
                    .Build(),
                ct);

            return watcher;
        }

        public int CountMatching(Func<byte[], bool> predicate) => _bodies.Count(b => Safe(predicate, b));

        public async Task<byte[]> WaitForOneAsync(TimeSpan timeout, CancellationToken ct)
            => await WaitForMatchAsync(_ => true, timeout, ct);

        public async Task<byte[]> WaitForMatchAsync(Func<byte[], bool> predicate, TimeSpan timeout, CancellationToken ct)
        {
            var deadline = DateTime.UtcNow + timeout;
            while (true)
            {
                if (_bodies.FirstOrDefault(b => Safe(predicate, b)) is { } match)
                {
                    return match;
                }

                if (DateTime.UtcNow >= deadline)
                {
                    throw new TimeoutException(
                        $"No matching STATE message within {timeout} ({_bodies.Count} seen, none matching).");
                }

                await Task.Delay(100, ct);
            }
        }

        private static bool Safe(Func<byte[], bool> predicate, byte[] body)
        {
            try { return predicate(body); }
            catch { return false; }
        }

        public async ValueTask DisposeAsync()
        {
            try
            {
                if (_client.IsConnected)
                {
                    await _client.DisconnectAsync();
                }
            }
            catch
            {
                // Teardown only.
            }
            finally
            {
                _client.Dispose();
            }
        }
    }
}