lmxopcua/tests/Core/ZB.MOM.WW.OtOpcUa.Commons.Tests/Browsing/DiscoveryDriverBrowserConcurrencyTests.cs

using Microsoft.Extensions.Logging.Abstractions;
using Shouldly;
using ZB.MOM.WW.OtOpcUa.Commons.Browsing;
using ZB.MOM.WW.OtOpcUa.Core.Abstractions;
using Xunit;

namespace ZB.MOM.WW.OtOpcUa.Commons.Tests.Browsing;

public class DiscoveryDriverBrowserConcurrencyTests
{
    private static readonly TimeSpan LongOpen = TimeSpan.FromSeconds(30);
    private static readonly TimeSpan ShortShutdown = TimeSpan.FromMilliseconds(200);
    private static readonly TimeSpan Wait = TimeSpan.FromSeconds(5);

    private static DiscoveryDriverBrowser Browser(IDriverFactory factory) =>
        new(factory, NullLogger<DiscoveryDriverBrowser>.Instance, LongOpen, ShortShutdown);

    /// <summary>A driver gate: signals when Initialize is entered, blocks until released.</summary>
    private sealed class Gate
    {
        public readonly SemaphoreSlim Entered = new(0);
        public readonly TaskCompletionSource Release = new(TaskCreationOptions.RunContinuationsAsynchronously);
    }

    private static FakeDiscoveryDriver Gated(Gate gate) => new()
    {
        OnInitialize = async ct =>
        {
            gate.Entered.Release();
            await gate.Release.Task.WaitAsync(ct);
        },
    };

    [Fact]
    public async Task SameKey_TwoConcurrentOpens_OneCapture_DistinctTokens()
    {
        var gate = new Gate();
        var factory = FakeDriverFactory.Of(() => Gated(gate));
        var browser = Browser(factory);

        var open1 = browser.OpenAsync("Fake", """{"h":"1"}""", CancellationToken.None);
        var open2 = browser.OpenAsync("Fake", """{"h":"1"}""", CancellationToken.None);

        await gate.Entered.WaitAsync(Wait, TestContext.Current.CancellationToken);          // one driver entered Initialize
        factory.CreateCalls.ShouldBe(1);             // coalesced — only one driver constructed
        gate.Release.SetResult();

        var s1 = await open1;
        var s2 = await open2;
        s1.Token.ShouldNotBe(s2.Token);              // distinct sessions...
        (await s1.RootAsync(CancellationToken.None)).Count
            .ShouldBe((await s2.RootAsync(CancellationToken.None)).Count); // ...over the same tree
        ((FakeDiscoveryDriver)factory.Created[0]).InitializeCalls.ShouldBe(1);
    }

    [Fact]
    public async Task SameKey_DisposingOneSession_LeavesOtherBrowsable()
    {
        var gate = new Gate();
        var browser = Browser(FakeDriverFactory.Of(() => Gated(gate)));

        var open1 = browser.OpenAsync("Fake", """{"h":"1"}""", CancellationToken.None);
        var open2 = browser.OpenAsync("Fake", """{"h":"1"}""", CancellationToken.None);
        await gate.Entered.WaitAsync(Wait, TestContext.Current.CancellationToken);
        gate.Release.SetResult();

        var s1 = await open1;
        var s2 = await open2;
        await s1.DisposeAsync();
        (await s2.RootAsync(CancellationToken.None)).ShouldNotBeEmpty(); // shared immutable tree survives
    }

    [Fact]
    public async Task SameKey_FailedCapture_FailsBothAwaiters()
    {
        var factory = FakeDriverFactory.Of(() => new FakeDiscoveryDriver
        {
            OnDiscover = (_, _) => throw new InvalidOperationException("shared-boom"),
        });
        var browser = Browser(factory);

        var open1 = browser.OpenAsync("Fake", """{"h":"1"}""", CancellationToken.None);
        var open2 = browser.OpenAsync("Fake", """{"h":"1"}""", CancellationToken.None);

        (await Should.ThrowAsync<InvalidOperationException>(() => open1)).Message.ShouldBe("shared-boom");
        (await Should.ThrowAsync<InvalidOperationException>(() => open2)).Message.ShouldBe("shared-boom");
    }

    [Fact]
    public async Task SameKey_OpenAfterCompletion_RunsFreshCapture()
    {
        var factory = FakeDriverFactory.Of(() => new FakeDiscoveryDriver()); // ungated, completes fast
        var browser = Browser(factory);

        await browser.OpenAsync("Fake", """{"h":"1"}""", CancellationToken.None);
        await browser.OpenAsync("Fake", """{"h":"1"}""", CancellationToken.None);

        factory.CreateCalls.ShouldBe(2); // in-flight-only map: second open is a fresh capture (Refresh path)
    }

    [Fact]
    public async Task DistinctKeys_CapLimitsInflightCaptures()
    {
        const int cap = DiscoveryDriverBrowser.MaxConcurrentCaptures; // 4
        const int total = cap + 2;                                    // 6
        var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        var entered = 0;
        var factory = new FakeDriverFactory((_, _, _) => new FakeDiscoveryDriver
        {
            OnInitialize = async ct => { Interlocked.Increment(ref entered); await release.Task.WaitAsync(ct); },
        });
        var browser = Browser(factory);

        var opens = Enumerable.Range(0, total)
            .Select(i => browser.OpenAsync("Fake", $$"""{"i":{{i}}}""", CancellationToken.None))
            .ToArray();

        // Give the semaphore time to admit exactly `cap` captures into Initialize.
        await Task.Delay(TimeSpan.FromMilliseconds(500), TestContext.Current.CancellationToken);
        Volatile.Read(ref entered).ShouldBe(cap); // at most 4 in flight; the other 2 queue on the semaphore

        release.SetResult();
        var sessions = await Task.WhenAll(opens);
        sessions.Length.ShouldBe(total);
        factory.CreateCalls.ShouldBe(total); // all eventually run once a slot frees
    }

    [Fact]
    public async Task QueuedOpen_ObservesCallerCancellation()
    {
        const int cap = DiscoveryDriverBrowser.MaxConcurrentCaptures; // 4
        var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        var factory = new FakeDriverFactory((_, _, _) => new FakeDiscoveryDriver
        {
            OnInitialize = async ct => { await release.Task.WaitAsync(ct); },
        });
        var browser = Browser(factory);

        // Fill all slots with distinct-key captures.
        var inflight = Enumerable.Range(0, cap)
            .Select(i => browser.OpenAsync("Fake", $$"""{"i":{{i}}}""", CancellationToken.None))
            .ToArray();
        await Task.Delay(TimeSpan.FromMilliseconds(300), TestContext.Current.CancellationToken);

        // A queued open with a cancelable token — it waits for a slot, then the caller cancels.
        using var cts = new CancellationTokenSource();
        var queued = browser.OpenAsync("Fake", """{"i":"queued"}""", cts.Token);
        cts.Cancel();
        await Should.ThrowAsync<OperationCanceledException>(() => queued);

        // The in-flight captures are unaffected.
        release.SetResult();
        var done = await Task.WhenAll(inflight);
        done.Length.ShouldBe(cap);
    }

    [Fact]
    public async Task CoalescedAwaiter_CancellationDoesNotCancelSharedCapture()
    {
        var gate = new Gate();
        var factory = FakeDriverFactory.Of(() => Gated(gate));
        var browser = Browser(factory);

        using var cts1 = new CancellationTokenSource();
        var open1 = browser.OpenAsync("Fake", """{"h":"1"}""", cts1.Token);
        var open2 = browser.OpenAsync("Fake", """{"h":"1"}""", CancellationToken.None);

        await gate.Entered.WaitAsync(Wait, TestContext.Current.CancellationToken);
        cts1.Cancel();                                    // first caller abandons
        await Should.ThrowAsync<OperationCanceledException>(() => open1);

        gate.Release.SetResult();                         // shared capture proceeds regardless
        var s2 = await open2;
        (await s2.RootAsync(CancellationToken.None)).ShouldNotBeEmpty();
        ((FakeDiscoveryDriver)factory.Created[0]).InitializeCalls.ShouldBe(1);
    }
}