natsdotnet/tests/NATS.Server.Mqtt.Tests/Mqtt/MqttFlowControllerTests.cs

// Go reference: server/mqtt.go — mqttMaxAckPending, flow control logic.

using NATS.Server.Mqtt;
using Shouldly;

namespace NATS.Server.Mqtt.Tests.Mqtt;

public sealed class MqttFlowControllerTests
{
    // 1. TryAcquire succeeds when under limit
    [Fact]
    public async Task TryAcquire_succeeds_when_under_limit()
    {
        using var fc = new MqttFlowController(defaultMaxAckPending: 1024);

        var result = await fc.TryAcquireAsync("sub-1");

        result.ShouldBeTrue();
    }

    // 2. TryAcquire fails when at limit
    [Fact]
    public async Task TryAcquire_fails_when_at_limit()
    {
        using var fc = new MqttFlowController(defaultMaxAckPending: 1);

        var first = await fc.TryAcquireAsync("sub-1");
        var second = await fc.TryAcquireAsync("sub-1");

        first.ShouldBeTrue();
        second.ShouldBeFalse();
    }

    // 3. Release allows next acquire
    [Fact]
    public async Task Release_allows_next_acquire()
    {
        using var fc = new MqttFlowController(defaultMaxAckPending: 1);

        var first = await fc.TryAcquireAsync("sub-1");
        first.ShouldBeTrue();

        // At limit — second should fail
        var atLimit = await fc.TryAcquireAsync("sub-1");
        atLimit.ShouldBeFalse();

        fc.Release("sub-1");

        // After release a slot is available again
        var afterRelease = await fc.TryAcquireAsync("sub-1");
        afterRelease.ShouldBeTrue();
    }

    // 4. GetPendingCount tracks pending
    [Fact]
    public async Task GetPendingCount_tracks_pending()
    {
        using var fc = new MqttFlowController(defaultMaxAckPending: 10);

        await fc.AcquireAsync("sub-1");
        await fc.AcquireAsync("sub-1");
        await fc.AcquireAsync("sub-1");

        fc.GetPendingCount("sub-1").ShouldBe(3);
    }

    // 5. GetPendingCount decrements on release
    [Fact]
    public async Task GetPendingCount_decrements_on_release()
    {
        using var fc = new MqttFlowController(defaultMaxAckPending: 10);

        await fc.AcquireAsync("sub-1");
        await fc.AcquireAsync("sub-1");
        await fc.AcquireAsync("sub-1");

        fc.Release("sub-1");

        fc.GetPendingCount("sub-1").ShouldBe(2);
    }

    // 6. GetPendingCount returns zero for unknown subscription
    [Fact]
    public void GetPendingCount_zero_for_unknown()
    {
        using var fc = new MqttFlowController();

        fc.GetPendingCount("does-not-exist").ShouldBe(0);
    }

    // 7. RemoveSubscription cleans up
    [Fact]
    public async Task RemoveSubscription_cleans_up()
    {
        using var fc = new MqttFlowController(defaultMaxAckPending: 10);

        await fc.AcquireAsync("sub-1");
        fc.SubscriptionCount.ShouldBe(1);

        fc.RemoveSubscription("sub-1");

        fc.SubscriptionCount.ShouldBe(0);
    }

    // 8. SubscriptionCount tracks independent subscriptions
    [Fact]
    public async Task SubscriptionCount_tracks_subscriptions()
    {
        using var fc = new MqttFlowController(defaultMaxAckPending: 10);

        await fc.AcquireAsync("sub-a");
        await fc.AcquireAsync("sub-b");
        await fc.AcquireAsync("sub-c");

        fc.SubscriptionCount.ShouldBe(3);
    }

    // 9. DefaultMaxAckPending can be updated via UpdateLimit
    [Fact]
    public void DefaultMaxAckPending_can_be_updated()
    {
        using var fc = new MqttFlowController(defaultMaxAckPending: 1024);
        fc.DefaultMaxAckPending.ShouldBe(1024);

        fc.UpdateLimit(512);

        fc.DefaultMaxAckPending.ShouldBe(512);
    }

    // 10. Dispose cleans up all subscriptions
    [Fact]
    public async Task Dispose_cleans_up_all()
    {
        var fc = new MqttFlowController(defaultMaxAckPending: 10);

        await fc.AcquireAsync("sub-x");
        await fc.AcquireAsync("sub-y");
        await fc.AcquireAsync("sub-z");

        fc.SubscriptionCount.ShouldBe(3);

        fc.Dispose();

        fc.SubscriptionCount.ShouldBe(0);
    }
}