lmxopcua/tests/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Modbus.Tests/ModbusHealthContractTests.cs

using System.Text.Json;
using Microsoft.Extensions.Logging;
using Shouldly;
using Xunit;
using ZB.MOM.WW.OtOpcUa.Core.Abstractions;

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

/// <summary>
///     The driver health contract (review 07 H4 + H7, and the lifecycle/fail-fast rules every driver lane
///     must meet): a device REFUSAL is per-tag Bad with no Degraded and proves the device answered; a comms
///     fault degrades and preserves LastSuccessfulRead; health is one write per batch; a comms timeout on
///     one unit fails the rest of that unit's batch fast; writes map transport failures to
///     BadCommunicationError and value problems to BadOutOfRange / BadTypeMismatch.
/// </summary>
[Trait("Category", "Unit")]
public sealed class ModbusHealthContractTests
{
    private const uint Good = 0u;
    private const uint BadOutOfRange = 0x803C0000u;
    private const uint BadDeviceFailure = 0x808B0000u;
    private const uint BadCommunicationError = 0x80050000u;
    private const uint BadTypeMismatch = 0x80740000u;
    private const uint BadNodeIdUnknown = 0x80340000u;

    private static ModbusDriverOptions Opts(params ModbusTagDefinition[] tags) => new()
    {
        Host = "f", RawTags = ModbusRawTags.Entries(tags), Probe = new ModbusProbeOptions { Enabled = false },
    };

    private static ModbusTagDefinition Hr(string name, ushort addr, byte? unit = null, ModbusDataType type = ModbusDataType.Int16) =>
        new(name, ModbusRegion.HoldingRegisters, addr, type, UnitId: unit);

    private static Task<byte[]> Refuse(byte code, byte[] pdu) =>
        Task.FromException<byte[]>(new ModbusException(pdu[0], code, $"exception {code}"));

    private static ushort Addr(byte[] pdu) => (ushort)((pdu[1] << 8) | pdu[2]);

    private static async Task<(ModbusDriver Driver, ScriptedTransport Transport)> StartAsync(
        ModbusDriverOptions opts, ILogger<ModbusDriver>? logger = null)
    {
        var t = new ScriptedTransport();
        var d = new ModbusDriver(opts, "m1", _ => t, logger);
        await d.InitializeAsync("{}", TestContext.Current.CancellationToken);
        return (d, t);
    }

    // ---------------- H4: refusals ----------------

    [Fact]
    public async Task All_tags_refused_keeps_Healthy_and_advances_LastSuccessfulRead()
    {
        var (drv, t) = await StartAsync(Opts(Hr("A", 10), Hr("B", 11)));
        await using var _ = drv;
        t.Handler = (_, pdu, _) => Refuse(0x02, pdu);
        var before = drv.GetHealth().LastSuccessfulRead;

        var r = await drv.ReadAsync(["A", "B"], TestContext.Current.CancellationToken);

        r.ShouldAllBe(s => s.StatusCode == BadOutOfRange);
        var h = drv.GetHealth();
        h.State.ShouldBe(DriverState.Healthy, "a refusal means the device answered — no Degraded, no watchdog fodder");
        h.LastSuccessfulRead.ShouldNotBeNull();
        h.LastSuccessfulRead!.Value.ShouldBeGreaterThanOrEqualTo(before!.Value);
    }

    [Fact]
    public async Task A_refusal_after_a_comms_fault_restores_Healthy()
    {
        var (drv, t) = await StartAsync(Opts(Hr("A", 10)));
        await using var _ = drv;
        t.Handler = (_, _, _) => Task.FromException<byte[]>(new IOException("socket dropped"));
        await drv.ReadAsync(["A"], TestContext.Current.CancellationToken);
        var degraded = drv.GetHealth();
        degraded.State.ShouldBe(DriverState.Degraded);

        t.Handler = (_, pdu, _) => Refuse(0x02, pdu);
        await drv.ReadAsync(["A"], TestContext.Current.CancellationToken);
        var h = drv.GetHealth();
        h.State.ShouldBe(DriverState.Healthy);
        h.LastSuccessfulRead!.Value.ShouldBeGreaterThanOrEqualTo(degraded.LastSuccessfulRead!.Value);
    }

    [Theory]
    [InlineData(0x04)]
    [InlineData(0x0B)]
    public async Task Non_refusal_exception_codes_still_degrade(byte code)
    {
        // Positive control for the refusal tests: a server-failure / gateway code is NOT a refusal.
        var (drv, t) = await StartAsync(Opts(Hr("A", 10)));
        await using var _ = drv;
        t.Handler = (_, pdu, _) => Refuse(code, pdu);
        await drv.ReadAsync(["A"], TestContext.Current.CancellationToken);
        drv.GetHealth().State.ShouldBe(DriverState.Degraded);
    }

    [Theory]
    [InlineData(false)]
    [InlineData(true)]
    public async Task Mixed_refused_and_good_batch_is_Healthy_regardless_of_tag_order(bool refusedFirst)
    {
        var (drv, t) = await StartAsync(Opts(Hr("Refused", 10), Hr("Fine", 50)));
        await using var _ = drv;
        t.Handler = (_, pdu, _) => Addr(pdu) == 10 ? Refuse(0x02, pdu) : null;
        string[] refs = refusedFirst ? ["Refused", "Fine"] : ["Fine", "Refused"];
        var r = await drv.ReadAsync(refs, TestContext.Current.CancellationToken);
        r.Count(s => s.StatusCode == Good).ShouldBe(1);
        drv.GetHealth().State.ShouldBe(DriverState.Healthy);
    }

    [Fact]
    public async Task A_refused_tag_warns_once_not_once_per_poll()
    {
        var logger = new ListLogger();
        var (drv, t) = await StartAsync(Opts(Hr("A", 10)), logger);
        await using var _ = drv;
        t.Handler = (_, pdu, _) => Refuse(0x02, pdu);
        for (var i = 0; i < 3; i++) await drv.ReadAsync(["A"], TestContext.Current.CancellationToken);
        logger.Warnings.Count(m => m.Contains("refused") && m.Contains("A")).ShouldBe(1);
    }

    // ---------------- fail fast per unit / endpoint ----------------

    [Fact]
    public async Task A_comms_timeout_on_one_unit_fails_the_rest_of_that_units_batch_fast()
    {
        var (drv, t) = await StartAsync(Opts(Hr("U2a", 10, 2), Hr("U2b", 11, 2), Hr("U2c", 12, 2), Hr("U1", 13, 1)));
        await using var _ = drv;
        t.Bank.HoldingRegisters[13] = 42;
        t.Handler = (unit, _, _) => unit == 2
            ? Task.FromException<byte[]>(new ModbusTransportDesyncException(
                ModbusTransportDesyncException.DesyncReason.Timeout, "timeout"))
            : null;

        var r = await drv.ReadAsync(["U2a", "U2b", "U2c", "U1"], TestContext.Current.CancellationToken);

        t.SendsTo(2).ShouldBe(1, "one timeout per unit per batch, not one per tag");
        r.Take(3).ShouldAllBe(s => s.StatusCode == BadCommunicationError);
        // Positive control: the healthy unit in the same batch is still read.
        r[3].StatusCode.ShouldBe(Good);
        r[3].Value.ShouldBe((short)42);
        drv.GetHealth().State.ShouldBe(DriverState.Degraded);
    }

    [Fact]
    public async Task An_unreachable_endpoint_fails_the_whole_batch_fast()
    {
        var (drv, t) = await StartAsync(Opts(Hr("A", 10, 1), Hr("B", 11, 2), Hr("C", 12, 3)));
        await using var _ = drv;
        t.Handler = (_, _, _) => Task.FromException<byte[]>(new ModbusTransportUnavailableException("throttled"));
        var r = await drv.ReadAsync(["A", "B", "C"], TestContext.Current.CancellationToken);
        t.TotalSends.ShouldBe(1);
        r.ShouldAllBe(s => s.StatusCode == BadCommunicationError);
    }

    // ---------------- H7: write status mapping + write-side health ----------------

    [Fact]
    public async Task Write_transport_failure_is_BadCommunicationError_and_degrades_the_write_side()
    {
        var (drv, t) = await StartAsync(Opts(Hr("A", 10)));
        await using var _ = drv;
        t.Handler = (_, pdu, _) => pdu[0] == 0x06 ? Task.FromException<byte[]>(new IOException("dropped")) : null;

        var w = await drv.WriteAsync([new WriteRequest("A", (short)1)], TestContext.Current.CancellationToken);
        w[0].StatusCode.ShouldBe(BadCommunicationError);
        drv.GetHealth().State.ShouldBe(DriverState.Degraded);

        // A good READ does not clear a write-side fault...
        (await drv.ReadAsync(["A"], TestContext.Current.CancellationToken))[0].StatusCode.ShouldBe(Good);
        drv.GetHealth().State.ShouldBe(DriverState.Degraded);

        // ...a good WRITE does.
        t.Handler = null;
        (await drv.WriteAsync([new WriteRequest("A", (short)2)], TestContext.Current.CancellationToken))[0].StatusCode.ShouldBe(Good);
        drv.GetHealth().State.ShouldBe(DriverState.Healthy);
    }

    [Fact]
    public async Task A_write_refusal_counts_as_answered_and_clears_a_write_side_fault()
    {
        var (drv, t) = await StartAsync(Opts(Hr("A", 10)));
        await using var _ = drv;
        t.Handler = (_, pdu, _) => pdu[0] == 0x06 ? Task.FromException<byte[]>(new IOException("dropped")) : null;
        await drv.WriteAsync([new WriteRequest("A", (short)1)], TestContext.Current.CancellationToken);
        drv.GetHealth().State.ShouldBe(DriverState.Degraded);

        t.Handler = (_, pdu, _) => pdu[0] == 0x06 ? Refuse(0x03, pdu) : null;
        var w = await drv.WriteAsync([new WriteRequest("A", (short)1)], TestContext.Current.CancellationToken);
        w[0].StatusCode.ShouldBe(BadOutOfRange);
        drv.GetHealth().State.ShouldBe(DriverState.Healthy);
    }

    [Fact]
    public async Task Write_value_problems_map_to_OutOfRange_and_TypeMismatch_without_touching_health()
    {
        var arr = new ModbusTagDefinition("Arr", ModbusRegion.HoldingRegisters, 20, ModbusDataType.Int16, ArrayCount: 3);
        var (drv, t) = await StartAsync(Opts(Hr("Bcd", 10, type: ModbusDataType.Bcd16), Hr("I16", 11), arr));
        await using var _ = drv;

        var w = await drv.WriteAsync([
            new WriteRequest("Bcd", 12345),             // > 4 BCD digits
            new WriteRequest("I16", 70000),             // > Int16
            new WriteRequest("I16", "not a number"),
            new WriteRequest("Arr", new short[] { 1, 2 }), // length mismatch
            new WriteRequest("Arr", new int[] { 1, 2, 3 }), // CLR list of another element type converts (ScadaBridge shape)
        ], TestContext.Current.CancellationToken);

        w[0].StatusCode.ShouldBe(BadOutOfRange);
        w[1].StatusCode.ShouldBe(BadOutOfRange);
        w[2].StatusCode.ShouldBe(BadTypeMismatch);
        w[3].StatusCode.ShouldBe(BadOutOfRange);
        w[4].StatusCode.ShouldBe(Good);
        t.Bank.HoldingRegisters[22].ShouldBe((ushort)3);
        drv.GetHealth().State.ShouldBe(DriverState.Healthy);
    }

    // ---------------- lifecycle generation ----------------

    [Fact]
    public async Task A_superseded_initialize_never_publishes_over_the_shutdown_that_overtook_it()
    {
        var t = new ScriptedTransport { ConnectGate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously) };
        var drv = new ModbusDriver(Opts(Hr("A", 10)), "m1", _ => t) { LifecycleGateTimeout = TimeSpan.FromSeconds(1) };
        var init = drv.InitializeAsync("{}", TestContext.Current.CancellationToken);
        await t.ConnectEntered.Task.WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);

        // The init is wedged in its connect; shutdown waits out the gate timeout and tears down anyway.
        await drv.ShutdownAsync(TestContext.Current.CancellationToken);
        drv.GetHealth().State.ShouldBe(DriverState.Unknown);

        t.ConnectGate.SetResult(); // the abandoned connect finally completes
        await Should.ThrowAsync<OperationCanceledException>(() => init);

        drv.GetHealth().State.ShouldBe(DriverState.Unknown, "the superseded init must not publish Healthy");
        t.DisposeCount.ShouldBe(1, "the superseded init releases the transport it built");
        await Should.ThrowAsync<InvalidOperationException>(() => drv.ReadAsync(["A"], TestContext.Current.CancellationToken));
    }

    [Fact]
    public async Task A_read_that_straddles_a_reinitialize_does_not_write_health_about_the_new_transport()
    {
        var first = new ScriptedTransport();
        var second = new ScriptedTransport();
        var release = new TaskCompletionSource<byte[]>(TaskCreationOptions.RunContinuationsAsynchronously);
        var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        first.Handler = (_, _, _) => { entered.TrySetResult(); return release.Task; };
        var calls = 0;
        var drv = new ModbusDriver(Opts(Hr("A", 10)), "m1", _ => Interlocked.Increment(ref calls) == 1 ? first : second);
        await using var _ = drv;
        await drv.InitializeAsync("{}", TestContext.Current.CancellationToken);

        var read = drv.ReadAsync(["A"], TestContext.Current.CancellationToken);
        await entered.Task.WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);

        await drv.InitializeAsync("{}", TestContext.Current.CancellationToken); // re-init onto `second`
        drv.GetHealth().State.ShouldBe(DriverState.Healthy);

        release.SetException(new IOException("old socket died")); // the straddling read now fails
        (await read)[0].StatusCode.ShouldBe(BadCommunicationError);
        drv.GetHealth().State.ShouldBe(DriverState.Healthy, "a read on the superseded transport must not degrade the new one");

        // Positive control: a failure on the CURRENT transport still degrades.
        second.Handler = (_, _, _) => Task.FromException<byte[]>(new IOException("new socket died"));
        await drv.ReadAsync(["A"], TestContext.Current.CancellationToken);
        drv.GetHealth().State.ShouldBe(DriverState.Degraded);
    }

    [Fact]
    public async Task Reinitialize_with_a_new_tag_set_drops_tags_the_new_config_no_longer_has()
    {
        var t = new ScriptedTransport();
        var drv = new ModbusDriver(Opts(Hr("Old", 10)), "m1", _ => t);
        await using var _ = drv;
        await drv.InitializeAsync("{}", TestContext.Current.CancellationToken);
        (await drv.ReadAsync(["Old"], TestContext.Current.CancellationToken))[0].StatusCode.ShouldBe(Good);

        var config = JsonSerializer.Serialize(new
        {
            host = "f",
            probe = new { enabled = false },
            rawTags = ModbusRawTags.Entries([Hr("New", 11)]),
        });
        await drv.InitializeAsync(config, TestContext.Current.CancellationToken);

        var r = await drv.ReadAsync(["Old", "New"], TestContext.Current.CancellationToken);
        r[0].StatusCode.ShouldBe(BadNodeIdUnknown, "a re-init publishes a fresh table, it does not merge into the old one");
        r[1].StatusCode.ShouldBe(Good);
    }

    // ---------------- complete teardown ----------------

    [Fact]
    public async Task Shutdown_joins_the_probe_loop_so_no_host_transition_fires_after_it_returns()
    {
        var t = new ScriptedTransport();
        var probeBlocked = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        var probes = 0;
        t.Handler = async (_, pdu, ct) =>
        {
            if (Interlocked.Increment(ref probes) == 1) return await t.Bank.SendAsync(1, pdu, ct); // first probe: Running
            probeBlocked.TrySetResult();
            try { await Task.Delay(Timeout.Infinite, ct); }
            catch (OperationCanceledException) { }
            throw new ObjectDisposedException("transport"); // a teardown-time failure, NOT an OCE
        };
        var opts = new ModbusDriverOptions
        {
            Host = "f", RawTags = ModbusRawTags.Entries([Hr("A", 10)]),
            Probe = new ModbusProbeOptions { Enabled = true, Interval = TimeSpan.FromMilliseconds(50), Timeout = TimeSpan.FromSeconds(30) },
        };
        var drv = new ModbusDriver(opts, "m1", _ => t);
        var running = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        var afterShutdown = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        var shutdownReturned = false;
        drv.OnHostStatusChanged += (_, e) =>
        {
            if (e.NewState == HostState.Running) running.TrySetResult();
            if (Volatile.Read(ref shutdownReturned)) afterShutdown.TrySetResult();
        };
        await drv.InitializeAsync("{}", TestContext.Current.CancellationToken);
        await running.Task.WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken); // positive control
        await probeBlocked.Task.WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);

        await drv.ShutdownAsync(TestContext.Current.CancellationToken);
        Volatile.Write(ref shutdownReturned, true);

        await Should.ThrowAsync<TimeoutException>(
            () => afterShutdown.Task.WaitAsync(TimeSpan.FromSeconds(1), TestContext.Current.CancellationToken));
    }

    private sealed class ListLogger : ILogger<ModbusDriver>
    {
        public System.Collections.Concurrent.ConcurrentQueue<string> Messages { get; } = new();
        public IEnumerable<string> Warnings => Messages.Where(m => m.StartsWith("Warning:", StringComparison.Ordinal));
        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)
            => Messages.Enqueue($"{logLevel}: {formatter(state, exception)}");
    }
}