lmxopcua/tests/Server/ZB.MOM.WW.OtOpcUa.ControlPlane.Tests/Communication/MeshPskProviderRefreshTests.cs

using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using Shouldly;
using Xunit;
using ZB.MOM.WW.OtOpcUa.Cluster;
using ZB.MOM.WW.OtOpcUa.ControlPlane.Communication;
using ZB.MOM.WW.Secrets.Abstractions;

namespace ZB.MOM.WW.OtOpcUa.ControlPlane.Tests.Communication;

/// <summary>
///     R03:3.4 — central invalidated every cached key on every contact refresh (default 60 s), so the
///     first inbound ack per cluster per minute blocked a gRPC thread on a synchronous, 3 s-bounded store
///     read and was denied past the bound. <see cref="MeshPskProvider.RefreshAsync"/> pre-resolves the
///     keys off the request path and swaps them in atomically.
/// </summary>
public sealed class MeshPskProviderRefreshTests
{
    private static MeshPskProvider Provider(ISecretResolver resolver) =>
        new(Options.Create(new MeshTransportOptions()), resolver, NullLogger<MeshPskProvider>.Instance);

    [Fact]
    public async Task After_a_refresh_the_lookup_is_served_without_touching_the_store()
    {
        var resolver = new GatedResolver { Value = "k-1" };
        var provider = Provider(resolver);

        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);
        resolver.Calls.ShouldBe(1);

        provider.TryGetPsk("SITE-A", out var psk).ShouldBeTrue();
        psk.ShouldBe("k-1");
        resolver.Calls.ShouldBe(1, "the request path did not consult the store");
        provider.Matches("SITE-A", "k-1").ShouldBeTrue();
    }

    /// <summary>
    ///     While a refresh is in flight the PREVIOUS key is still served (no invalidate-then-reread window);
    ///     once it completes the rotated key is served. Positive control on both sides of the swap.
    /// </summary>
    [Fact]
    public async Task A_rotated_key_is_swapped_in_atomically_with_the_old_key_served_until_then()
    {
        var resolver = new GatedResolver { Value = "old" };
        var provider = Provider(resolver);
        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);
        provider.TryGetPsk("SITE-A", out var before).ShouldBeTrue();
        before.ShouldBe("old");

        resolver.Value = "new";
        resolver.Hold();
        var refresh = provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);
        resolver.Entered.Wait(TimeSpan.FromSeconds(5)).ShouldBeTrue();

        // Mid-refresh: the old key, from memory — the store is NOT re-read on this path.
        var callsMidRefresh = resolver.Calls;
        provider.TryGetPsk("SITE-A", out var during).ShouldBeTrue();
        during.ShouldBe("old");
        resolver.Calls.ShouldBe(callsMidRefresh);

        resolver.Release();
        await refresh;
        provider.TryGetPsk("SITE-A", out var after).ShouldBeTrue();
        after.ShouldBe("new");
    }

    [Fact]
    public async Task A_cluster_whose_refresh_fails_keeps_its_previous_key()
    {
        var resolver = new GatedResolver { Value = "k-1" };
        var provider = Provider(resolver);
        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);

        resolver.Throw = new InvalidOperationException("store offline");
        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);

        provider.TryGetPsk("SITE-A", out var psk).ShouldBeTrue();
        psk.ShouldBe("k-1");
    }

    /// <summary>A cluster that was never refreshed still resolves lazily on the request path (unchanged behaviour).</summary>
    [Fact]
    public async Task A_cluster_outside_the_refresh_set_still_resolves_lazily()
    {
        var resolver = new GatedResolver { Value = "k-b" };
        var provider = Provider(resolver);
        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);

        provider.TryGetPsk("SITE-B", out var psk).ShouldBeTrue();
        psk.ShouldBe("k-b");
        resolver.Calls.ShouldBe(2);
    }

    [Fact]
    public async Task InvalidateAll_drops_refreshed_keys_too()
    {
        var resolver = new GatedResolver { Value = "k-1" };
        var provider = Provider(resolver);
        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);

        provider.InvalidateAll();
        resolver.Value = "k-2";
        provider.TryGetPsk("SITE-A", out var psk).ShouldBeTrue();
        psk.ShouldBe("k-2");
    }

    private sealed class GatedResolver : ISecretResolver
    {
        private int _calls;
        private volatile TaskCompletionSource? _gate;

        public string? Value { get; set; }
        public Exception? Throw { get; set; }
        public int Calls => Volatile.Read(ref _calls);
        public ManualResetEventSlim Entered { get; } = new();

        public void Hold() => _gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        public void Release() => _gate?.TrySetResult();

        public async Task<string?> GetAsync(SecretName name, CancellationToken ct)
        {
            Interlocked.Increment(ref _calls);
            Entered.Set();
            if (_gate is { } gate) await gate.Task.WaitAsync(ct);
            if (Throw is not null) throw Throw;
            return Value;
        }
    }
}