natsdotnet/tests/NATS.Server.Transport.Tests/TlsRateLimiterTests.cs

using NATS.Server.TestUtilities;
using NATS.Server.Tls;

namespace NATS.Server.Transport.Tests;

public class TlsRateLimiterTests
{
    [Fact]
    public async Task Rate_limiter_allows_configured_tokens_per_second()
    {
        using var limiter = new TlsRateLimiter(5);

        // Should allow 5 tokens immediately
        for (int i = 0; i < 5; i++)
        {
            using var cts = new CancellationTokenSource(100);
            await limiter.WaitAsync(cts.Token); // Should not throw
        }

        // 6th token should block (no refill yet)
        using var blockCts = new CancellationTokenSource(200);
        var blocked = false;
        try
        {
            await limiter.WaitAsync(blockCts.Token);
        }
        catch (OperationCanceledException)
        {
            blocked = true;
        }
        blocked.ShouldBeTrue("6th token should be blocked before refill");
    }

    [Fact]
    public async Task Rate_limiter_refills_after_one_second()
    {
        using var limiter = new TlsRateLimiter(2);

        // Consume all tokens
        await limiter.WaitAsync(CancellationToken.None);
        await limiter.WaitAsync(CancellationToken.None);

        // Wait for refill (rate limiter refills tokens after 1 second)
        await PollHelper.YieldForAsync(1200);

        // Should have tokens again
        using var cts = new CancellationTokenSource(200);
        await limiter.WaitAsync(cts.Token); // Should not throw
    }
}