natsdotnet/tests/NATS.Server.Tests/Auth/AccountLimitsTests.cs

// Tests for per-account JetStream resource limits.
// Go reference: accounts_test.go TestAccountLimits, TestJetStreamLimits.

using NATS.Server.Auth;

namespace NATS.Server.Tests.Auth;

public class AccountLimitsTests
{
    [Fact]
    public void TryReserveConsumer_UnderLimit_ReturnsTrue()
    {
        var account = new Account("test")
        {
            JetStreamLimits = new AccountLimits { MaxConsumers = 3 },
        };

        account.TryReserveConsumer().ShouldBeTrue();
        account.TryReserveConsumer().ShouldBeTrue();
        account.TryReserveConsumer().ShouldBeTrue();
        account.ConsumerCount.ShouldBe(3);
    }

    [Fact]
    public void TryReserveConsumer_AtLimit_ReturnsFalse()
    {
        var account = new Account("test")
        {
            JetStreamLimits = new AccountLimits { MaxConsumers = 2 },
        };

        account.TryReserveConsumer().ShouldBeTrue();
        account.TryReserveConsumer().ShouldBeTrue();
        account.TryReserveConsumer().ShouldBeFalse();
        account.ConsumerCount.ShouldBe(2);
    }

    [Fact]
    public void ReleaseConsumer_DecrementsCount()
    {
        var account = new Account("test")
        {
            JetStreamLimits = new AccountLimits { MaxConsumers = 2 },
        };

        account.TryReserveConsumer().ShouldBeTrue();
        account.TryReserveConsumer().ShouldBeTrue();
        account.ConsumerCount.ShouldBe(2);

        account.ReleaseConsumer();
        account.ConsumerCount.ShouldBe(1);

        // Now we can reserve again
        account.TryReserveConsumer().ShouldBeTrue();
        account.ConsumerCount.ShouldBe(2);
    }

    [Fact]
    public void TrackStorageDelta_UnderLimit_ReturnsTrue()
    {
        var account = new Account("test")
        {
            JetStreamLimits = new AccountLimits { MaxStorage = 1000 },
        };

        account.TrackStorageDelta(500).ShouldBeTrue();
        account.StorageUsed.ShouldBe(500);

        account.TrackStorageDelta(400).ShouldBeTrue();
        account.StorageUsed.ShouldBe(900);
    }

    [Fact]
    public void TrackStorageDelta_ExceedsLimit_ReturnsFalse()
    {
        var account = new Account("test")
        {
            JetStreamLimits = new AccountLimits { MaxStorage = 1000 },
        };

        account.TrackStorageDelta(800).ShouldBeTrue();
        account.TrackStorageDelta(300).ShouldBeFalse(); // 800 + 300 = 1100 > 1000
        account.StorageUsed.ShouldBe(800); // unchanged
    }

    [Fact]
    public void TrackStorageDelta_NegativeDelta_ReducesUsage()
    {
        var account = new Account("test")
        {
            JetStreamLimits = new AccountLimits { MaxStorage = 1000 },
        };

        account.TrackStorageDelta(800).ShouldBeTrue();
        account.TrackStorageDelta(-300).ShouldBeTrue(); // negative always succeeds
        account.StorageUsed.ShouldBe(500);

        // Now we have room again
        account.TrackStorageDelta(400).ShouldBeTrue();
        account.StorageUsed.ShouldBe(900);
    }

    [Fact]
    public void MaxStorage_Zero_Unlimited()
    {
        var account = new Account("test")
        {
            JetStreamLimits = new AccountLimits { MaxStorage = 0 }, // unlimited
        };

        // Should accept any amount
        account.TrackStorageDelta(long.MaxValue / 2).ShouldBeTrue();
        account.StorageUsed.ShouldBe(long.MaxValue / 2);
    }

    [Fact]
    public void Limits_DefaultValues_AllUnlimited()
    {
        var limits = AccountLimits.Unlimited;

        limits.MaxStorage.ShouldBe(0);
        limits.MaxStreams.ShouldBe(0);
        limits.MaxConsumers.ShouldBe(0);
        limits.MaxAckPending.ShouldBe(0);
        limits.MaxMemoryStorage.ShouldBe(0);
        limits.MaxDiskStorage.ShouldBe(0);

        // Account defaults to unlimited
        var account = new Account("test");
        account.JetStreamLimits.ShouldBe(AccountLimits.Unlimited);
    }

    [Fact]
    public void TryReserveStream_WithLimits_RespectsNewLimits()
    {
        // JetStreamLimits.MaxStreams should take precedence over MaxJetStreamStreams
        var account = new Account("test")
        {
            MaxJetStreamStreams = 10, // legacy field
            JetStreamLimits = new AccountLimits { MaxStreams = 2 }, // new limit overrides
        };

        account.TryReserveStream().ShouldBeTrue();
        account.TryReserveStream().ShouldBeTrue();
        account.TryReserveStream().ShouldBeFalse(); // limited to 2 by JetStreamLimits
        account.JetStreamStreamCount.ShouldBe(2);
    }

    [Fact]
    public void EvictOldestClient_WhenMaxConnectionsExceeded()
    {
        var account = new Account("test")
        {
            MaxConnections = 2,
        };

        account.AddClient(1).ShouldBeTrue();
        account.AddClient(2).ShouldBeTrue();
        account.AddClient(3).ShouldBeFalse(); // at limit
        account.ClientCount.ShouldBe(2);

        // Remove oldest, then new one can connect
        account.RemoveClient(1);
        account.ClientCount.ShouldBe(1);

        account.AddClient(3).ShouldBeTrue();
        account.ClientCount.ShouldBe(2);
    }
}