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

// Reference: golang/nats-server/server/route.go:533-545 — computeRoutePoolIdx
// Tests for account-based route pool index computation and message routing.

using System.Text;
using Microsoft.Extensions.Logging.Abstractions;
using NATS.Server.Configuration;
using NATS.Server.Routes;

namespace NATS.Server.Clustering.Tests.Routes;

/// <summary>
/// Tests for route pool accounting per account, matching Go's
/// computeRoutePoolIdx behavior (route.go:533-545).
/// </summary>
public class RoutePoolAccountTests
{
    [Fact]
    public void ComputeRoutePoolIdx_SinglePool_AlwaysReturnsZero()
    {
        RouteManager.ComputeRoutePoolIdx(1, "account-A").ShouldBe(0);
        RouteManager.ComputeRoutePoolIdx(1, "account-B").ShouldBe(0);
        RouteManager.ComputeRoutePoolIdx(1, "$G").ShouldBe(0);
        RouteManager.ComputeRoutePoolIdx(0, "anything").ShouldBe(0);
    }

    [Fact]
    public void ComputeRoutePoolIdx_DeterministicForSameAccount()
    {
        const int poolSize = 5;
        const string account = "my-test-account";

        var first = RouteManager.ComputeRoutePoolIdx(poolSize, account);
        var second = RouteManager.ComputeRoutePoolIdx(poolSize, account);
        var third = RouteManager.ComputeRoutePoolIdx(poolSize, account);

        first.ShouldBe(second);
        second.ShouldBe(third);
        first.ShouldBeGreaterThanOrEqualTo(0);
        first.ShouldBeLessThan(poolSize);
    }

    [Fact]
    public void ComputeRoutePoolIdx_DistributesAcrossPool()
    {
        const int poolSize = 3;
        var usedIndices = new HashSet<int>();

        for (var i = 0; i < 100; i++)
        {
            var idx = RouteManager.ComputeRoutePoolIdx(poolSize, $"account-{i}");
            idx.ShouldBeGreaterThanOrEqualTo(0);
            idx.ShouldBeLessThan(poolSize);
            usedIndices.Add(idx);
        }

        usedIndices.Count.ShouldBe(poolSize);
    }

    [Fact]
    public void ComputeRoutePoolIdx_EmptyAccount_ReturnsValid()
    {
        const int poolSize = 4;
        var idx = RouteManager.ComputeRoutePoolIdx(poolSize, string.Empty);
        idx.ShouldBeGreaterThanOrEqualTo(0);
        idx.ShouldBeLessThan(poolSize);
    }

    [Fact]
    public void ComputeRoutePoolIdx_DefaultGlobalAccount_ReturnsValid()
    {
        const int poolSize = 3;
        var idx = RouteManager.ComputeRoutePoolIdx(poolSize, "$G");
        idx.ShouldBeGreaterThanOrEqualTo(0);
        idx.ShouldBeLessThan(poolSize);

        var idx2 = RouteManager.ComputeRoutePoolIdx(poolSize, "$G");
        idx.ShouldBe(idx2);
    }

    [Fact]
    public async Task ForwardRoutedMessage_UsesCorrectPoolConnection()
    {
        var clusterName = Guid.NewGuid().ToString("N");

        var optsA = new NatsOptions
        {
            Host = "127.0.0.1",
            Port = 0,
            Cluster = new ClusterOptions
            {
                Name = clusterName,
                Host = "127.0.0.1",
                Port = 0,
                PoolSize = 1,
                Routes = [],
            },
        };

        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,
            Cluster = new ClusterOptions
            {
                Name = clusterName,
                Host = "127.0.0.1",
                Port = 0,
                PoolSize = 1,
                Routes = [$"127.0.0.1:{optsA.Cluster.Port}"],
            },
        };

        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 &&
               (Interlocked.Read(ref serverA.Stats.Routes) == 0 ||
                Interlocked.Read(ref serverB.Stats.Routes) == 0))
        {
            await Task.Delay(50, timeout.Token);
        }

        Interlocked.Read(ref serverA.Stats.Routes).ShouldBeGreaterThan(0);

        var payload = Encoding.UTF8.GetBytes("hello");
        await serverA.RouteManager!.ForwardRoutedMessageAsync(
            "$G", "test.subject", null, payload, CancellationToken.None);

        var poolIdx = RouteManager.ComputeRoutePoolIdx(1, "$G");
        poolIdx.ShouldBe(0);

        await ctsA.CancelAsync();
        await ctsB.CancelAsync();
        serverA.Dispose();
        serverB.Dispose();
        ctsA.Dispose();
        ctsB.Dispose();
    }
}