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

using System.Collections.Concurrent;
using System.Reflection;
using Shouldly;
using Xunit;
using ZB.MOM.WW.OtOpcUa.Core.Abstractions;
using ZB.MOM.WW.OtOpcUa.Driver.Modbus;

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

/// <summary>
///     Regression coverage for Driver.Modbus findings -002 (Reinitialize state hygiene),
///     -003 (_health volatile-write ordering), -004 (DisposeAsync teardown parity), and
///     -005 (malformed/short response PDU handling). All four resolved fixes need a
///     unit test alongside them per Driver.Modbus-012.
/// </summary>
[Trait("Category", "Unit")]
public sealed class ModbusLifecycleHygieneTests
{
    private sealed class FakeTransport : IModbusTransport
    {
        public readonly ushort[] HoldingRegisters = new ushort[256];
        public int ConnectCount;
        public int DisposeCount;
        public int SendCount;

        public Task ConnectAsync(CancellationToken ct) { Interlocked.Increment(ref ConnectCount); return Task.CompletedTask; }

        public Task<byte[]> SendAsync(byte unitId, byte[] pdu, CancellationToken ct)
        {
            Interlocked.Increment(ref SendCount);
            var fc = pdu[0];
            switch (fc)
            {
                case 0x03:
                case 0x04:
                {
                    var addr = (ushort)((pdu[1] << 8) | pdu[2]);
                    var qty = (ushort)((pdu[3] << 8) | pdu[4]);
                    var resp = new byte[2 + qty * 2];
                    resp[0] = fc;
                    resp[1] = (byte)(qty * 2);
                    for (var i = 0; i < qty; i++)
                    {
                        resp[2 + i * 2] = (byte)(HoldingRegisters[addr + i] >> 8);
                        resp[3 + i * 2] = (byte)(HoldingRegisters[addr + i] & 0xFF);
                    }
                    return Task.FromResult(resp);
                }
                case 0x06:
                {
                    var addr = (ushort)((pdu[1] << 8) | pdu[2]);
                    HoldingRegisters[addr] = (ushort)((pdu[3] << 8) | pdu[4]);
                    return Task.FromResult(pdu); // FC06 echoes the request
                }
                case 0x10:
                {
                    var addr = (ushort)((pdu[1] << 8) | pdu[2]);
                    var qty = (ushort)((pdu[3] << 8) | pdu[4]);
                    for (var i = 0; i < qty; i++)
                        HoldingRegisters[addr + i] = (ushort)((pdu[6 + i * 2] << 8) | pdu[7 + i * 2]);
                    return Task.FromResult(new byte[] { 0x10, pdu[1], pdu[2], pdu[3], pdu[4] });
                }
                default:
                    return Task.FromException<byte[]>(new NotSupportedException($"fc={fc}"));
            }
        }

        public ValueTask DisposeAsync() { Interlocked.Increment(ref DisposeCount); return ValueTask.CompletedTask; }
    }

    /// <summary>
    ///     Returns a snapshot of the driver's private <c>_tagsByName</c> dictionary so the
    ///     hygiene tests can confirm the cache is empty after teardown.
    /// </summary>
    private static System.Collections.IDictionary GetTagsByName(ModbusDriver drv) =>
        (System.Collections.IDictionary)typeof(ModbusDriver)
            .GetField("_tagsByName", BindingFlags.NonPublic | BindingFlags.Instance)!
            .GetValue(drv)!;

    // -------------------- Finding -002 / -012 (2) --------------------

    [Fact]
    public async Task Reinitialize_clears_stale_tagsByName_entries()
    {
        // Re-initializing with a different options instance would leak stale entries before
        // the fix. We simulate by inspecting _tagsByName after a Shutdown — it must be empty
        // so InitializeAsync repopulates from a clean slate.
        var fake = new FakeTransport();
        var opts = new ModbusDriverOptions
        {
            Host = "fake",
            Tags = [new ModbusTagDefinition("A", ModbusRegion.HoldingRegisters, 0, ModbusDataType.Int16)],
        };
        var drv = new ModbusDriver(opts, "modbus-1", _ => fake);
        await drv.InitializeAsync("{}", CancellationToken.None);

        GetTagsByName(drv).Count.ShouldBe(1);
        await drv.ShutdownAsync(CancellationToken.None);
        GetTagsByName(drv).Count.ShouldBe(0, "Shutdown must clear the tag cache so the next Initialize starts clean");
    }

    [Fact]
    public async Task Reinitialize_clears_lastPublished_and_lastWritten_caches()
    {
        var fake = new FakeTransport();
        var opts = new ModbusDriverOptions
        {
            Host = "fake",
            WriteOnChangeOnly = true,
            Tags = [new ModbusTagDefinition("A", ModbusRegion.HoldingRegisters, 0, ModbusDataType.Int16,
                Deadband: 1.0)],
        };
        var drv = new ModbusDriver(opts, "modbus-1", _ => fake);
        await drv.InitializeAsync("{}", CancellationToken.None);

        var lastPublished = (System.Collections.IDictionary)typeof(ModbusDriver)
            .GetField("_lastPublishedByRef", BindingFlags.NonPublic | BindingFlags.Instance)!
            .GetValue(drv)!;
        var lastWritten = (System.Collections.IDictionary)typeof(ModbusDriver)
            .GetField("_lastWrittenByRef", BindingFlags.NonPublic | BindingFlags.Instance)!
            .GetValue(drv)!;

        // Seed both caches via a write (lastWritten) and a publish through ShouldPublish (lastPublished).
        await drv.WriteAsync([new WriteRequest("A", (short)5)], CancellationToken.None);
        lastWritten.Count.ShouldBe(1);

        // Reach ShouldPublish directly through a subscription so the deadband cache fills.
        fake.HoldingRegisters[0] = 5;
        var handle = await drv.SubscribeAsync(["A"], TimeSpan.FromMilliseconds(100), CancellationToken.None);
        var deadline = DateTime.UtcNow.AddSeconds(2);
        while (lastPublished.Count == 0 && DateTime.UtcNow < deadline) await Task.Delay(25);
        lastPublished.Count.ShouldBe(1);
        await drv.UnsubscribeAsync(handle, CancellationToken.None);

        await drv.ShutdownAsync(CancellationToken.None);
        lastPublished.Count.ShouldBe(0, "Shutdown must clear the deadband cache");
        lastWritten.Count.ShouldBe(0, "Shutdown must clear the write-suppression cache");
    }

    // -------------------- Finding -004 / -012 (4) --------------------

    [Fact]
    public async Task DisposeAsync_without_explicit_Shutdown_tears_down_probe_loop_and_transport()
    {
        var fake = new FakeTransport();
        var opts = new ModbusDriverOptions
        {
            Host = "fake",
            Probe = new ModbusProbeOptions
            {
                Enabled = true,
                Interval = TimeSpan.FromMilliseconds(50),
                Timeout = TimeSpan.FromSeconds(1),
            },
            // Re-probe loop also opted in so DisposeAsync exercises both CTS cancellations.
            AutoProhibitReprobeInterval = TimeSpan.FromMilliseconds(50),
            Tags = [new ModbusTagDefinition("A", ModbusRegion.HoldingRegisters, 0, ModbusDataType.Int16)],
        };
        var drv = new ModbusDriver(opts, "modbus-1", _ => fake);
        await drv.InitializeAsync("{}", CancellationToken.None);

        // Let the probe + re-probe loops spin a few iterations.
        await Task.Delay(200);
        var sendsAtDispose = Interlocked.CompareExchange(ref fake.SendCount, 0, 0);
        sendsAtDispose.ShouldBeGreaterThan(0, "background probe loop should have issued at least one send");

        // Skip ShutdownAsync — exercise the await-using path that previously leaked.
        await drv.DisposeAsync();

        // Transport must have been disposed exactly once and the background loops stop scheduling
        // new sends. Tolerate at most one in-flight send straddling the cancel.
        fake.DisposeCount.ShouldBe(1);
        var sendsAfterDispose = Interlocked.CompareExchange(ref fake.SendCount, 0, 0);
        await Task.Delay(300);
        var sendsAtRest = Interlocked.CompareExchange(ref fake.SendCount, 0, 0);
        (sendsAtRest - sendsAfterDispose).ShouldBeLessThanOrEqualTo(1, "background loops must stop after DisposeAsync");
    }

    [Fact]
    public async Task DisposeAsync_disposes_the_pollEngine_so_subscriptions_stop()
    {
        var fake = new FakeTransport();
        var opts = new ModbusDriverOptions
        {
            Host = "fake",
            Probe = new ModbusProbeOptions { Enabled = false },
            Tags = [new ModbusTagDefinition("A", ModbusRegion.HoldingRegisters, 0, ModbusDataType.Int16)],
        };
        var drv = new ModbusDriver(opts, "modbus-1", _ => fake);
        await drv.InitializeAsync("{}", CancellationToken.None);

        // Spin up a polled subscription; the PollGroupEngine schedules a background Task that
        // will keep issuing SendAsync until either Unsubscribe or DisposeAsync stops it.
        var handle = await drv.SubscribeAsync(["A"], TimeSpan.FromMilliseconds(100), CancellationToken.None);
        await Task.Delay(250);
        var beforeDispose = Interlocked.CompareExchange(ref fake.SendCount, 0, 0);
        beforeDispose.ShouldBeGreaterThan(0);

        // No ShutdownAsync — DisposeAsync must also tear down the poll engine.
        await drv.DisposeAsync();

        var atDispose = Interlocked.CompareExchange(ref fake.SendCount, 0, 0);
        await Task.Delay(400);
        var atRest = Interlocked.CompareExchange(ref fake.SendCount, 0, 0);

        (atRest - atDispose).ShouldBeLessThanOrEqualTo(1,
            "DisposeAsync must dispose the PollGroupEngine so its background Task stops, not just the transport");
    }

    // -------------------- Finding -005 / -012 (3) --------------------

    /// <summary>
    ///     Transport that returns a structurally-broken response for FC03/FC04 — too short to
    ///     hold the declared byte-count. Pre-fix the driver dereferenced <c>resp[1]</c> and then
    ///     ran <c>Buffer.BlockCopy(resp, 2, ..., resp[1])</c> which threw <c>ArgumentException</c>
    ///     (out-of-range). Post-fix the driver throws <c>InvalidDataException</c> which the
    ///     <c>ReadAsync</c> catch-all maps to <see cref="BadCommunicationError"/>.
    /// </summary>
    private sealed class TruncatingTransport : IModbusTransport
    {
        /// <summary>How many bytes to return — anything &lt; 2 + bytecount is malformed.</summary>
        public int ResponseBytes { get; set; } = 1; // just the fc byte, no bytecount

        public Task ConnectAsync(CancellationToken ct) => Task.CompletedTask;
        public Task<byte[]> SendAsync(byte unitId, byte[] pdu, CancellationToken ct)
        {
            var resp = new byte[ResponseBytes];
            if (ResponseBytes >= 1) resp[0] = pdu[0];
            if (ResponseBytes >= 2) resp[1] = 4; // claim 4 bytes of payload but provide none
            return Task.FromResult(resp);
        }
        public ValueTask DisposeAsync() => ValueTask.CompletedTask;
    }

    [Fact]
    public async Task Short_response_PDU_surfaces_as_BadCommunicationError_not_an_IndexOutOfRangeException()
    {
        var fake = new TruncatingTransport { ResponseBytes = 1 };
        var opts = new ModbusDriverOptions
        {
            Host = "fake",
            Tags = [new ModbusTagDefinition("Level", ModbusRegion.HoldingRegisters, 0, ModbusDataType.Int16)],
        };
        var drv = new ModbusDriver(opts, "modbus-1", _ => fake);
        await drv.InitializeAsync("{}", CancellationToken.None);

        var r = await drv.ReadAsync(["Level"], CancellationToken.None);
        r[0].StatusCode.ShouldBe(0x80050000u, "BadCommunicationError = a clean transport-layer fault");
        r[0].Value.ShouldBeNull();
    }

    [Fact]
    public async Task Response_payload_truncated_below_declared_byteCount_surfaces_as_BadCommunicationError()
    {
        // Header says "4 bytes follow" but the message is only 3 bytes total — pre-fix the
        // Buffer.BlockCopy would throw ArgumentException.
        var fake = new TruncatingTransport { ResponseBytes = 3 };
        var opts = new ModbusDriverOptions
        {
            Host = "fake",
            Tags = [new ModbusTagDefinition("Level", ModbusRegion.HoldingRegisters, 0, ModbusDataType.Int16)],
        };
        var drv = new ModbusDriver(opts, "modbus-1", _ => fake);
        await drv.InitializeAsync("{}", CancellationToken.None);

        var r = await drv.ReadAsync(["Level"], CancellationToken.None);
        r[0].StatusCode.ShouldBe(0x80050000u);
    }

    [Fact]
    public void DecodeBitArray_rejects_an_empty_bitmap_with_InvalidDataException()
    {
        var decode = typeof(ModbusDriver).GetMethod(
            "DecodeBitArray", BindingFlags.NonPublic | BindingFlags.Static)!;
        // We can't invoke through reflection because ReadOnlySpan<byte> isn't representable in
        // object-array invocation parameters. Instead, exercise the path through ReadAsync with
        // a bit-region tag and a transport that returns a zero-byte-count response.
        var fake = new EmptyBitTransport();
        var opts = new ModbusDriverOptions
        {
            Host = "fake",
            Tags = [new ModbusTagDefinition("Coil", ModbusRegion.Coils, 0, ModbusDataType.Bool)],
        };
        var drv = new ModbusDriver(opts, "modbus-1", _ => fake);
        drv.InitializeAsync("{}", CancellationToken.None).GetAwaiter().GetResult();

        var r = drv.ReadAsync(["Coil"], CancellationToken.None).GetAwaiter().GetResult();
        // The empty-bitmap guard surfaces via the BadCommunicationError catch-all.
        r[0].StatusCode.ShouldBe(0x80050000u);
    }

    /// <summary>
    ///     Coil-bank transport that returns <c>[fc][bytecount=0]</c> — a response with a
    ///     declared zero-byte payload. Pre-fix <c>DecodeBitArray</c> indexed into the empty
    ///     bitmap and threw <c>IndexOutOfRangeException</c>.
    /// </summary>
    private sealed class EmptyBitTransport : IModbusTransport
    {
        public Task ConnectAsync(CancellationToken ct) => Task.CompletedTask;
        public Task<byte[]> SendAsync(byte unitId, byte[] pdu, CancellationToken ct)
            => Task.FromResult(new byte[] { pdu[0], 0 });
        public ValueTask DisposeAsync() => ValueTask.CompletedTask;
    }

    // -------------------- Finding -003 (volatile _health) --------------------

    /// <summary>
    ///     The <c>_health</c> field is read by <c>GetHealth()</c> and written by every read /
    ///     write / probe path. The fix uses <c>Volatile.Read</c>/<c>Volatile.Write</c> to give
    ///     <c>GetHealth()</c> a defined ordering guarantee. We verify that under concurrent
    ///     pressure <c>GetHealth()</c> never returns a half-constructed value (it's a sealed
    ///     record so reference-assignment atomicity already prevents tearing; the test guards
    ///     against future regressions to a struct-typed health surface).
    /// </summary>
    [Fact]
    public async Task GetHealth_under_concurrent_pressure_always_returns_a_complete_snapshot()
    {
        var fake = new FakeTransport();
        var opts = new ModbusDriverOptions
        {
            Host = "fake",
            Tags = [new ModbusTagDefinition("A", ModbusRegion.HoldingRegisters, 0, ModbusDataType.Int16)],
        };
        var drv = new ModbusDriver(opts, "modbus-1", _ => fake);
        await drv.InitializeAsync("{}", CancellationToken.None);

        // Two writer loops and one reader loop — 250ms of churn.
        var cts = new CancellationTokenSource(TimeSpan.FromMilliseconds(250));
        var faults = new ConcurrentQueue<Exception>();
        var writer = Task.Run(async () =>
        {
            try { while (!cts.IsCancellationRequested) await drv.ReadAsync(["A"], CancellationToken.None); }
            catch (Exception ex) { faults.Enqueue(ex); }
        });
        var reader = Task.Run(() =>
        {
            try
            {
                while (!cts.IsCancellationRequested)
                {
                    var h = drv.GetHealth();
                    // State must be one of the enum values; LastSuccessfulRead can be null or a real time;
                    // the record constructor enforces no field is wholly garbage.
                    h.State.ShouldBeOneOf(DriverState.Unknown, DriverState.Initializing, DriverState.Healthy, DriverState.Degraded, DriverState.Faulted);
                }
            }
            catch (Exception ex) { faults.Enqueue(ex); }
        });

        await Task.WhenAll(writer, reader);
        faults.ShouldBeEmpty();
    }
}