natsdotnet/tests/NATS.Server.JetStream.Tests/JetStream/Consumers/ClusterPendingRequestTests.cs

using NATS.Server.JetStream.Consumers;
using NATS.Server.JetStream.Cluster;
using Shouldly;

namespace NATS.Server.JetStream.Tests.JetStream.Consumers;

/// <summary>
/// Tests for cluster-aware pending pull request tracking in PullConsumerEngine.
/// Go reference: consumer.go proposeWaitingRequest / waitingRequestsPending — cluster-wide
/// pending pull request coordination via the consumer RAFT group.
/// golang/nats-server/server/consumer.go proposeWaitingRequest
/// </summary>
public class ClusterPendingRequestTests
{
    // ---------------------------------------------------------------
    // ProposeWaitingRequest
    // ---------------------------------------------------------------

    [Fact]
    public void ProposeWaitingRequest_with_quorum_returns_true()
    {
        // Go: consumer.go proposeWaitingRequest — only propose when quorum available.
        var engine = new PullConsumerEngine();
        var group = new RaftGroup
        {
            Name = "test-group",
            Peers = ["peer-1", "peer-2", "peer-3"],
        };
        var request = new PullWaitingRequest { Batch = 10, Reply = "reply.test.1" };

        var result = engine.ProposeWaitingRequest(request, group);

        result.ShouldBeTrue();
    }

    [Fact]
    public void ProposeWaitingRequest_without_quorum_returns_false()
    {
        // Go: consumer.go proposeWaitingRequest — no quorum (0 peers means quorum = 1, but 0 < 1).
        var engine = new PullConsumerEngine();
        var group = new RaftGroup
        {
            Name = "empty-group",
            Peers = [],
        };
        var request = new PullWaitingRequest { Batch = 5, Reply = "reply.noquorum" };

        var result = engine.ProposeWaitingRequest(request, group);

        result.ShouldBeFalse();
    }

    [Fact]
    public void ProposeWaitingRequest_registers_in_cluster_pending()
    {
        // Go: consumer.go — after a successful proposal, the request must appear in the
        // cluster pending map so it can be fulfilled or expired.
        var engine = new PullConsumerEngine();
        var group = new RaftGroup
        {
            Name = "test-group",
            Peers = ["peer-1", "peer-2", "peer-3"],
        };
        var request = new PullWaitingRequest { Batch = 4, Reply = "reply.reg" };

        engine.ProposeWaitingRequest(request, group);

        var pending = engine.GetClusterPendingRequests();
        pending.ShouldContain(r => r.Reply == "reply.reg");
    }

    [Fact]
    public void Multiple_proposals_tracked_independently()
    {
        // Go: consumer.go — each reply subject is an independent pending slot;
        // proposals with different reply subjects must not overwrite each other.
        var engine = new PullConsumerEngine();
        var group = new RaftGroup
        {
            Name = "test-group",
            Peers = ["peer-1", "peer-2", "peer-3"],
        };

        engine.ProposeWaitingRequest(new PullWaitingRequest { Batch = 1, Reply = "reply.A" }, group);
        engine.ProposeWaitingRequest(new PullWaitingRequest { Batch = 2, Reply = "reply.B" }, group);
        engine.ProposeWaitingRequest(new PullWaitingRequest { Batch = 3, Reply = "reply.C" }, group);

        engine.ClusterPendingCount.ShouldBe(3);
        var pending = engine.GetClusterPendingRequests();
        pending.ShouldContain(r => r.Reply == "reply.A" && r.Batch == 1);
        pending.ShouldContain(r => r.Reply == "reply.B" && r.Batch == 2);
        pending.ShouldContain(r => r.Reply == "reply.C" && r.Batch == 3);
    }

    // ---------------------------------------------------------------
    // ClusterPendingCount
    // ---------------------------------------------------------------

    [Fact]
    public void ClusterPendingCount_tracks_pending_requests()
    {
        // Go: consumer.go — ClusterPendingCount reflects the current size of the pending map.
        var engine = new PullConsumerEngine();
        engine.ClusterPendingCount.ShouldBe(0);

        engine.RegisterClusterPending(new PullWaitingRequest { Batch = 1, Reply = "r1" });
        engine.RegisterClusterPending(new PullWaitingRequest { Batch = 2, Reply = "r2" });
        engine.RegisterClusterPending(new PullWaitingRequest { Batch = 3, Reply = "r3" });

        engine.ClusterPendingCount.ShouldBe(3);
    }

    [Fact]
    public void ClusterPendingCount_decrements_on_remove()
    {
        // Go: consumer.go — removing a request via reply subject decrements the pending count.
        var engine = new PullConsumerEngine();
        engine.RegisterClusterPending(new PullWaitingRequest { Batch = 5, Reply = "decrement.reply" });
        engine.ClusterPendingCount.ShouldBe(1);

        engine.RemoveClusterPending("decrement.reply");

        engine.ClusterPendingCount.ShouldBe(0);
    }

    // ---------------------------------------------------------------
    // RegisterClusterPending
    // ---------------------------------------------------------------

    [Fact]
    public void RegisterClusterPending_adds_request_by_reply()
    {
        // Go: consumer.go — pending requests are keyed by reply subject for O(1) lookup.
        var engine = new PullConsumerEngine();
        var request = new PullWaitingRequest { Batch = 7, Reply = "register.reply.subject" };

        engine.RegisterClusterPending(request);

        var retrieved = engine.RemoveClusterPending("register.reply.subject");
        retrieved.ShouldNotBeNull();
        retrieved.Batch.ShouldBe(7);
        retrieved.Reply.ShouldBe("register.reply.subject");
    }

    // ---------------------------------------------------------------
    // RemoveClusterPending
    // ---------------------------------------------------------------

    [Fact]
    public void RemoveClusterPending_returns_and_removes_request()
    {
        // Go: consumer.go — RemoveClusterPending both returns the request and removes it
        // from the map so it is not fulfilled twice.
        var engine = new PullConsumerEngine();
        engine.RegisterClusterPending(new PullWaitingRequest { Batch = 3, Reply = "remove.me" });

        var removed = engine.RemoveClusterPending("remove.me");

        removed.ShouldNotBeNull();
        removed.Reply.ShouldBe("remove.me");
        engine.ClusterPendingCount.ShouldBe(0);
        // Second removal should return null — the entry is gone.
        engine.RemoveClusterPending("remove.me").ShouldBeNull();
    }

    [Fact]
    public void RemoveClusterPending_returns_null_for_unknown()
    {
        // Go: consumer.go — attempting to remove an unknown reply subject is a no-op.
        var engine = new PullConsumerEngine();

        var result = engine.RemoveClusterPending("does.not.exist");

        result.ShouldBeNull();
    }

    // ---------------------------------------------------------------
    // GetClusterPendingRequests
    // ---------------------------------------------------------------

    [Fact]
    public void GetClusterPendingRequests_returns_all_pending()
    {
        // Go: consumer.go — GetClusterPendingRequests is used for expiry sweeps and
        // diagnostics; it must return every currently pending request.
        var engine = new PullConsumerEngine();
        engine.RegisterClusterPending(new PullWaitingRequest { Batch = 1, Reply = "bulk.a" });
        engine.RegisterClusterPending(new PullWaitingRequest { Batch = 2, Reply = "bulk.b" });
        engine.RegisterClusterPending(new PullWaitingRequest { Batch = 3, Reply = "bulk.c" });

        var all = engine.GetClusterPendingRequests();

        all.Count.ShouldBe(3);
        all.Select(r => r.Reply).ShouldContain("bulk.a");
        all.Select(r => r.Reply).ShouldContain("bulk.b");
        all.Select(r => r.Reply).ShouldContain("bulk.c");
    }
}