lmxopcua/tests/Core/ZB.MOM.WW.OtOpcUa.Core.Tests/DriverHostTests.cs

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

namespace ZB.MOM.WW.OtOpcUa.Core.Tests;

[Trait("Category", "Unit")]
public sealed class DriverHostTests
{
    private sealed class StubDriver(string id, bool failInit = false) : IDriver
    {
        public string DriverInstanceId { get; } = id;
        public string DriverType => "Stub";
        public bool Initialized { get; private set; }
        public bool ShutDown { get; private set; }

        public Task InitializeAsync(string _, CancellationToken ct)
        {
            if (failInit) throw new InvalidOperationException("boom");
            Initialized = true;
            return Task.CompletedTask;
        }

        public Task ReinitializeAsync(string _, CancellationToken ct) => Task.CompletedTask;
        public Task ShutdownAsync(CancellationToken ct) { ShutDown = true; return Task.CompletedTask; }
        public DriverHealth GetHealth() =>
            new(Initialized ? DriverState.Healthy : DriverState.Unknown, null, null);
        public long GetMemoryFootprint() => 0;
        public Task FlushOptionalCachesAsync(CancellationToken ct) => Task.CompletedTask;
    }

    [Fact]
    public async Task Register_initializes_driver_and_tracks_health()
    {
        await using var host = new DriverHost();
        var driver = new StubDriver("d-1");

        await host.RegisterAsync(driver, "{}", CancellationToken.None);

        host.RegisteredDriverIds.ShouldContain("d-1");
        driver.Initialized.ShouldBeTrue();
        host.GetHealth("d-1")!.State.ShouldBe(DriverState.Healthy);
    }

    [Fact]
    public async Task Register_rethrows_init_failure_but_keeps_driver_registered()
    {
        await using var host = new DriverHost();
        var driver = new StubDriver("d-bad", failInit: true);

        await Should.ThrowAsync<InvalidOperationException>(() =>
            host.RegisterAsync(driver, "{}", CancellationToken.None));

        host.RegisteredDriverIds.ShouldContain("d-bad");
    }

    [Fact]
    public async Task Duplicate_registration_throws()
    {
        await using var host = new DriverHost();
        await host.RegisterAsync(new StubDriver("d-1"), "{}", CancellationToken.None);

        await Should.ThrowAsync<InvalidOperationException>(() =>
            host.RegisterAsync(new StubDriver("d-1"), "{}", CancellationToken.None));
    }

    [Fact]
    public async Task Unregister_shuts_down_and_removes()
    {
        await using var host = new DriverHost();
        var driver = new StubDriver("d-1");
        await host.RegisterAsync(driver, "{}", CancellationToken.None);

        await host.UnregisterAsync("d-1", CancellationToken.None);

        host.RegisteredDriverIds.ShouldNotContain("d-1");
        driver.ShutDown.ShouldBeTrue();
    }

    /// <summary>
    ///     Core-004 regression — DriverHost is a library type whose async calls must use
    ///     ConfigureAwait(false) to match the convention used by CapabilityInvoker /
    ///     AlarmSurfaceInvoker. Asserts the awaited driver call does not post its
    ///     continuation back to a captured SynchronizationContext.
    ///     The driver awaits an unsettled TaskCompletionSource so it does not introduce its
    ///     own capture — only DriverHost's await of the returned Task can drive a post.
    /// </summary>
    [Fact]
    public async Task RegisterAsync_Does_Not_Capture_SynchronizationContext()
    {
        var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        var driver = new TcsDriver("d-cfg-1", tcs);
        var ctx = new TrackingSynchronizationContext();

        // Run the DriverHost call on a dedicated thread that has our tracking SyncContext installed.
        var workerCtx = await RunOnContextAsync(ctx, async () =>
        {
            var host = new DriverHost();
            var registerTask = host.RegisterAsync(driver, "{}", CancellationToken.None);
            // Complete the driver's InitializeAsync from a background thread so DriverHost's
            // await must resume via the captured context if ConfigureAwait(false) was missing.
            _ = Task.Run(() => tcs.SetResult());
            await registerTask.ConfigureAwait(false);
            await host.DisposeAsync().ConfigureAwait(false);
        });

        workerCtx.PostCount.ShouldBe(0,
            "RegisterAsync's awaited driver call must use ConfigureAwait(false) so the continuation does not post back to the captured context");
    }

    [Fact]
    public async Task UnregisterAsync_Does_Not_Capture_SynchronizationContext()
    {
        var initTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        initTcs.SetResult();
        var shutdownTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        var driver = new TcsDriver("d-cfg-2", initTcs, shutdownTcs);
        var ctx = new TrackingSynchronizationContext();

        var workerCtx = await RunOnContextAsync(ctx, async () =>
        {
            var host = new DriverHost();
            await host.RegisterAsync(driver, "{}", CancellationToken.None).ConfigureAwait(false);

            // After RegisterAsync we re-enter the context. Reset the post counter so we only
            // observe UnregisterAsync's behaviour from here on.
            ((TrackingSynchronizationContext)SynchronizationContext.Current!).Reset();

            var task = host.UnregisterAsync("d-cfg-2", CancellationToken.None);
            _ = Task.Run(() => shutdownTcs.SetResult());
            await task.ConfigureAwait(false);
        });

        workerCtx.PostCount.ShouldBe(0,
            "UnregisterAsync's awaited shutdown call must use ConfigureAwait(false)");
    }

    [Fact]
    public async Task DisposeAsync_Does_Not_Capture_SynchronizationContext()
    {
        var initTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        initTcs.SetResult();
        var shutdownTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        var driver = new TcsDriver("d-cfg-3", initTcs, shutdownTcs);
        var ctx = new TrackingSynchronizationContext();

        var workerCtx = await RunOnContextAsync(ctx, async () =>
        {
            var host = new DriverHost();
            await host.RegisterAsync(driver, "{}", CancellationToken.None).ConfigureAwait(false);
            ((TrackingSynchronizationContext)SynchronizationContext.Current!).Reset();

            var task = host.DisposeAsync();
            _ = Task.Run(() => shutdownTcs.SetResult());
            await task.ConfigureAwait(false);
        });

        workerCtx.PostCount.ShouldBe(0,
            "DisposeAsync's awaited shutdown call must use ConfigureAwait(false)");
    }

    /// <summary>
    ///     Run <paramref name="body"/> on a dedicated thread with <paramref name="ctx"/>
    ///     installed as the current SynchronizationContext, and return <paramref name="ctx"/>
    ///     after the body completes. The dedicated thread guarantees that resuming via the
    ///     captured context observably routes through our Post hook (the ThreadPool would
    ///     otherwise clear the context on the resuming worker).
    /// </summary>
    private static Task<TrackingSynchronizationContext> RunOnContextAsync(TrackingSynchronizationContext ctx, Func<Task> body)
    {
        var done = new TaskCompletionSource<TrackingSynchronizationContext>(TaskCreationOptions.RunContinuationsAsynchronously);
        var t = new Thread(() =>
        {
            SynchronizationContext.SetSynchronizationContext(ctx);
            try
            {
                // Pump posted continuations until the body completes.
                var task = body();
                while (!task.IsCompleted)
                {
                    if (ctx.TryDequeue(out var work)) work();
                    else Thread.Sleep(1);
                }
                // Drain any tail continuations.
                while (ctx.TryDequeue(out var work)) work();
                task.GetAwaiter().GetResult();
                done.SetResult(ctx);
            }
            catch (Exception ex) { done.SetException(ex); }
        }) { IsBackground = true };
        t.Start();
        return done.Task;
    }

    /// <summary>Driver whose Initialize / Shutdown completions are caller-controlled via TCS.</summary>
    private sealed class TcsDriver(string id, TaskCompletionSource initTcs, TaskCompletionSource? shutdownTcs = null) : IDriver
    {
        public string DriverInstanceId { get; } = id;
        public string DriverType => "Tcs";

        public Task InitializeAsync(string _, CancellationToken ct) => initTcs.Task;
        public Task ReinitializeAsync(string _, CancellationToken ct) => Task.CompletedTask;
        public Task ShutdownAsync(CancellationToken ct) => (shutdownTcs ?? CompletedTcs).Task;
        public DriverHealth GetHealth() => new(DriverState.Healthy, null, null);
        public long GetMemoryFootprint() => 0;
        public Task FlushOptionalCachesAsync(CancellationToken ct) => Task.CompletedTask;

        private static readonly TaskCompletionSource CompletedTcs = MakeCompleted();
        private static TaskCompletionSource MakeCompleted()
        {
            var t = new TaskCompletionSource();
            t.SetResult();
            return t;
        }
    }

    /// <summary>SynchronizationContext that queues posts to a thread-safe work list and counts them.</summary>
    private sealed class TrackingSynchronizationContext : SynchronizationContext
    {
        private readonly System.Collections.Concurrent.ConcurrentQueue<Action> _queue = new();
        public int PostCount;
        public int SendCount;

        public override void Post(SendOrPostCallback d, object? state)
        {
            Interlocked.Increment(ref PostCount);
            _queue.Enqueue(() => d(state));
        }
        public override void Send(SendOrPostCallback d, object? state)
        {
            Interlocked.Increment(ref SendCount);
            d(state);
        }
        public bool TryDequeue(out Action work) => _queue.TryDequeue(out work!);
        public void Reset() { Interlocked.Exchange(ref PostCount, 0); Interlocked.Exchange(ref SendCount, 0); }
    }
}