natsdotnet/tests/NATS.Server.JetStream.Tests/JetStream/Cluster/JetStreamInflightTrackingTests.cs

using NATS.Server.JetStream.Cluster;

namespace NATS.Server.JetStream.Tests.JetStream.Cluster;

public class JetStreamInflightTrackingTests
{
    [Fact]
    public void TrackInflightStreamProposal_increments_ops()
    {
        var meta = new JetStreamMetaGroup(3);
        var sa = new StreamAssignment
        {
            StreamName = "inflight-1",
            Group = new RaftGroup { Name = "rg-inf", Peers = ["n1", "n2", "n3"] },
        };

        meta.TrackInflightStreamProposal("ACC", sa);
        meta.InflightStreamCount.ShouldBe(1);
        meta.IsStreamInflight("ACC", "inflight-1").ShouldBeTrue();
    }

    [Fact]
    public void RemoveInflightStreamProposal_clears_when_zero()
    {
        var meta = new JetStreamMetaGroup(3);
        var sa = new StreamAssignment
        {
            StreamName = "inflight-2",
            Group = new RaftGroup { Name = "rg-inf2", Peers = ["n1", "n2", "n3"] },
        };

        meta.TrackInflightStreamProposal("ACC", sa);
        meta.RemoveInflightStreamProposal("ACC", "inflight-2");
        meta.IsStreamInflight("ACC", "inflight-2").ShouldBeFalse();
    }

    [Fact]
    public void Duplicate_proposal_increments_ops_count()
    {
        var meta = new JetStreamMetaGroup(3);
        var sa = new StreamAssignment
        {
            StreamName = "dup-stream",
            Group = new RaftGroup { Name = "rg-dup", Peers = ["n1", "n2", "n3"] },
        };

        meta.TrackInflightStreamProposal("ACC", sa);
        meta.TrackInflightStreamProposal("ACC", sa);
        meta.InflightStreamCount.ShouldBe(1); // still one unique stream

        // Need two removes to fully clear
        meta.RemoveInflightStreamProposal("ACC", "dup-stream");
        meta.IsStreamInflight("ACC", "dup-stream").ShouldBeTrue(); // ops > 0
        meta.RemoveInflightStreamProposal("ACC", "dup-stream");
        meta.IsStreamInflight("ACC", "dup-stream").ShouldBeFalse();
    }

    [Fact]
    public void IsStreamInflight_returns_false_for_unknown_account()
    {
        var meta = new JetStreamMetaGroup(3);
        meta.IsStreamInflight("UNKNOWN", "no-stream").ShouldBeFalse();
    }

    [Fact]
    public void TrackInflightConsumerProposal_tracks_by_account()
    {
        var meta = new JetStreamMetaGroup(3);
        meta.TrackInflightConsumerProposal("ACC", "stream1", "consumer1");

        meta.InflightConsumerCount.ShouldBe(1);
        meta.IsConsumerInflight("ACC", "stream1", "consumer1").ShouldBeTrue();
    }

    [Fact]
    public void RemoveInflightConsumerProposal_clears_when_zero()
    {
        var meta = new JetStreamMetaGroup(3);
        meta.TrackInflightConsumerProposal("ACC", "stream1", "consumer1");
        meta.RemoveInflightConsumerProposal("ACC", "stream1", "consumer1");

        meta.IsConsumerInflight("ACC", "stream1", "consumer1").ShouldBeFalse();
    }

    [Fact]
    public void ClearAllInflight_removes_everything()
    {
        var meta = new JetStreamMetaGroup(3);
        var sa = new StreamAssignment
        {
            StreamName = "s1",
            Group = new RaftGroup { Name = "rg", Peers = ["n1", "n2", "n3"] },
        };

        meta.TrackInflightStreamProposal("ACC1", sa);
        meta.TrackInflightConsumerProposal("ACC2", "s2", "c1");

        meta.ClearAllInflight();

        meta.InflightStreamCount.ShouldBe(0);
        meta.InflightConsumerCount.ShouldBe(0);
    }

    [Fact]
    public void StepDown_clears_inflight()
    {
        var meta = new JetStreamMetaGroup(3);
        var sa = new StreamAssignment
        {
            StreamName = "s1",
            Group = new RaftGroup { Name = "rg", Peers = ["n1", "n2", "n3"] },
        };

        meta.TrackInflightStreamProposal("ACC", sa);
        meta.StepDown();

        meta.InflightStreamCount.ShouldBe(0);
    }

    [Fact]
    public void Multiple_accounts_tracked_independently()
    {
        var meta = new JetStreamMetaGroup(3);
        var sa1 = new StreamAssignment
        {
            StreamName = "s1",
            Group = new RaftGroup { Name = "rg1", Peers = ["n1", "n2", "n3"] },
        };
        var sa2 = new StreamAssignment
        {
            StreamName = "s1", // same stream name, different account
            Group = new RaftGroup { Name = "rg2", Peers = ["n1", "n2", "n3"] },
        };

        meta.TrackInflightStreamProposal("ACC1", sa1);
        meta.TrackInflightStreamProposal("ACC2", sa2);

        meta.InflightStreamCount.ShouldBe(2); // one per account
        meta.IsStreamInflight("ACC1", "s1").ShouldBeTrue();
        meta.IsStreamInflight("ACC2", "s1").ShouldBeTrue();

        meta.RemoveInflightStreamProposal("ACC1", "s1");
        meta.IsStreamInflight("ACC1", "s1").ShouldBeFalse();
        meta.IsStreamInflight("ACC2", "s1").ShouldBeTrue(); // still tracked
    }
}