natsdotnet/tests/NATS.Server.Clustering.Tests/Routes/RouteAccountScopedDeliveryTests.cs

using Microsoft.Extensions.Logging.Abstractions;
using NATS.Client.Core;
using NATS.Server.Auth;
using NATS.Server.Configuration;

namespace NATS.Server.Clustering.Tests.Routes;

public class RouteAccountScopedDeliveryTests
{
    [Fact]
    public async Task Remote_message_delivery_uses_target_account_sublist_not_global_sublist()
    {
        const string subject = "orders.created";
        await using var fixture = await RouteAccountDeliveryFixture.StartAsync();

        await using var remoteAccountA = await fixture.ConnectAsync(fixture.ServerB, "a_sub");
        await using var remoteAccountB = await fixture.ConnectAsync(fixture.ServerB, "b_sub");
        await using var publisher = await fixture.ConnectAsync(fixture.ServerA, "a_pub");

        await using var subA = await remoteAccountA.SubscribeCoreAsync<string>(subject);
        await using var subB = await remoteAccountB.SubscribeCoreAsync<string>(subject);
        await remoteAccountA.PingAsync();
        await remoteAccountB.PingAsync();
        await fixture.WaitForRemoteInterestOnServerAAsync("A", subject);

        await publisher.PublishAsync(subject, "from-route-a");

        using var receiveTimeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        var msgA = await subA.Msgs.ReadAsync(receiveTimeout.Token);
        msgA.Data.ShouldBe("from-route-a");

        using var leakTimeout = new CancellationTokenSource(TimeSpan.FromMilliseconds(500));
        await Should.ThrowAsync<OperationCanceledException>(async () =>
            await subB.Msgs.ReadAsync(leakTimeout.Token));
    }
}

internal sealed class RouteAccountDeliveryFixture : IAsyncDisposable
{
    private readonly CancellationTokenSource _ctsA;
    private readonly CancellationTokenSource _ctsB;

    private RouteAccountDeliveryFixture(NatsServer serverA, NatsServer serverB, CancellationTokenSource ctsA, CancellationTokenSource ctsB)
    {
        ServerA = serverA;
        ServerB = serverB;
        _ctsA = ctsA;
        _ctsB = ctsB;
    }

    public NatsServer ServerA { get; }
    public NatsServer ServerB { get; }

    public static async Task<RouteAccountDeliveryFixture> StartAsync()
    {
        var users = new User[]
        {
            new() { Username = "a_pub", Password = "pass", Account = "A" },
            new() { Username = "a_sub", Password = "pass", Account = "A" },
            new() { Username = "b_sub", Password = "pass", Account = "B" },
        };

        var optsA = new NatsOptions
        {
            Host = "127.0.0.1",
            Port = 0,
            Users = users,
            Cluster = new ClusterOptions
            {
                Name = Guid.NewGuid().ToString("N"),
                Host = "127.0.0.1",
                Port = 0,
            },
        };

        var serverA = new NatsServer(optsA, NullLoggerFactory.Instance);
        var ctsA = new CancellationTokenSource();
        _ = serverA.StartAsync(ctsA.Token);
        await serverA.WaitForReadyAsync();

        var optsB = new NatsOptions
        {
            Host = "127.0.0.1",
            Port = 0,
            Users = users,
            Cluster = new ClusterOptions
            {
                Name = Guid.NewGuid().ToString("N"),
                Host = "127.0.0.1",
                Port = 0,
                Routes = [serverA.ClusterListen!],
            },
        };

        var serverB = new NatsServer(optsB, NullLoggerFactory.Instance);
        var ctsB = new CancellationTokenSource();
        _ = serverB.StartAsync(ctsB.Token);
        await serverB.WaitForReadyAsync();

        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        while (!timeout.IsCancellationRequested && (serverA.Stats.Routes == 0 || serverB.Stats.Routes == 0))
            await Task.Delay(50, timeout.Token).ContinueWith(_ => { }, TaskScheduler.Default);

        return new RouteAccountDeliveryFixture(serverA, serverB, ctsA, ctsB);
    }

    public async Task<NatsConnection> ConnectAsync(NatsServer server, string username)
    {
        var connection = new NatsConnection(new NatsOpts
        {
            Url = $"nats://{username}:pass@127.0.0.1:{server.Port}",
        });
        await connection.ConnectAsync();
        return connection;
    }

    public async Task WaitForRemoteInterestOnServerAAsync(string account, string subject)
    {
        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        while (!timeout.IsCancellationRequested)
        {
            if (ServerA.HasRemoteInterest(account, subject))
                return;

            await Task.Delay(50, timeout.Token).ContinueWith(_ => { }, TaskScheduler.Default);
        }

        throw new TimeoutException($"Timed out waiting for remote interest {account}:{subject}.");
    }

    public async ValueTask DisposeAsync()
    {
        await _ctsA.CancelAsync();
        await _ctsB.CancelAsync();
        ServerA.Dispose();
        ServerB.Dispose();
        _ctsA.Dispose();
        _ctsB.Dispose();
    }
}