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

// Go parity: golang/nats-server/server/jetstream_cluster.go
// Covers: StreamReplicaGroup construction from StreamAssignment, per-stream RAFT apply
// logic (processStreamEntries), checkpoint/restore snapshot lifecycle, and commit/processed
// index tracking through the group facade.
using NATS.Server.JetStream.Cluster;
using NATS.Server.Raft;

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

/// <summary>
/// Tests for B10: per-stream RAFT apply logic added to StreamReplicaGroup.
/// Covers construction from StreamAssignment, apply loop, snapshot checkpoint/restore,
/// and the CommitIndex/ProcessedIndex/PendingCommits facade properties.
/// Go reference: jetstream_cluster.go processStreamAssignment, processStreamEntries.
/// </summary>
public class StreamReplicaGroupApplyTests
{
    // ---------------------------------------------------------------
    // Go: jetstream_cluster.go processStreamAssignment — builds per-stream raft group
    // ---------------------------------------------------------------

    [Fact]
    public void Construction_from_assignment_creates_correct_number_of_nodes()
    {
        var assignment = new StreamAssignment
        {
            StreamName = "ORDERS",
            Group = new RaftGroup
            {
                Name = "orders-raft",
                Peers = ["n1", "n2", "n3"],
            },
        };

        var group = new StreamReplicaGroup(assignment);

        group.Nodes.Count.ShouldBe(3);
        group.StreamName.ShouldBe("ORDERS");
        group.Assignment.ShouldNotBeNull();
        group.Assignment!.StreamName.ShouldBe("ORDERS");
    }

    [Fact]
    public void Construction_from_assignment_uses_peer_ids_as_node_ids()
    {
        var assignment = new StreamAssignment
        {
            StreamName = "EVENTS",
            Group = new RaftGroup
            {
                Name = "events-raft",
                Peers = ["peer-a", "peer-b", "peer-c"],
            },
        };

        var group = new StreamReplicaGroup(assignment);

        var nodeIds = group.Nodes.Select(n => n.Id).ToHashSet();
        nodeIds.ShouldContain("peer-a");
        nodeIds.ShouldContain("peer-b");
        nodeIds.ShouldContain("peer-c");
    }

    [Fact]
    public void Construction_from_assignment_elects_leader()
    {
        var assignment = new StreamAssignment
        {
            StreamName = "STREAM",
            Group = new RaftGroup
            {
                Name = "stream-raft",
                Peers = ["n1", "n2", "n3"],
            },
        };

        var group = new StreamReplicaGroup(assignment);

        group.Leader.ShouldNotBeNull();
        group.Leader.IsLeader.ShouldBeTrue();
    }

    [Fact]
    public void Construction_from_assignment_with_no_peers_creates_single_node()
    {
        var assignment = new StreamAssignment
        {
            StreamName = "SOLO",
            Group = new RaftGroup { Name = "solo-raft" },
        };

        var group = new StreamReplicaGroup(assignment);

        group.Nodes.Count.ShouldBe(1);
        group.Leader.IsLeader.ShouldBeTrue();
    }

    // ---------------------------------------------------------------
    // Go: raft.go:150-160 (applied/processed fields) — commit index on proposal
    // ---------------------------------------------------------------

    [Fact]
    public async Task ProposeAsync_through_group_increments_commit_index()
    {
        var group = new StreamReplicaGroup("TRACK", replicas: 3);
        group.CommitIndex.ShouldBe(0);

        await group.ProposeAsync("msg.1", default);

        group.CommitIndex.ShouldBe(1);
    }

    [Fact]
    public async Task Multiple_proposals_increment_commit_index_monotonically()
    {
        var group = new StreamReplicaGroup("MULTI", replicas: 3);

        await group.ProposeAsync("msg.1", default);
        await group.ProposeAsync("msg.2", default);
        await group.ProposeAsync("msg.3", default);

        group.CommitIndex.ShouldBe(3);
    }

    // ---------------------------------------------------------------
    // Go: jetstream_cluster.go processStreamEntries — apply loop
    // ---------------------------------------------------------------

    [Fact]
    public async Task ApplyCommittedEntriesAsync_processes_pending_entries()
    {
        var group = new StreamReplicaGroup("APPLY", replicas: 3);

        await group.ProposeAsync("store.msg.1", default);
        await group.ProposeAsync("store.msg.2", default);

        group.PendingCommits.ShouldBe(2);

        await group.ApplyCommittedEntriesAsync(default);

        group.PendingCommits.ShouldBe(0);
        group.ProcessedIndex.ShouldBe(2);
    }

    [Fact]
    public async Task ApplyCommittedEntriesAsync_marks_regular_entries_as_processed()
    {
        var group = new StreamReplicaGroup("MARK", replicas: 1);

        var idx = await group.ProposeAsync("data.record", default);

        group.ProcessedIndex.ShouldBe(0);

        await group.ApplyCommittedEntriesAsync(default);

        group.ProcessedIndex.ShouldBe(idx);
    }

    [Fact]
    public async Task ApplyCommittedEntriesAsync_on_empty_queue_is_noop()
    {
        var group = new StreamReplicaGroup("EMPTY", replicas: 3);

        // No proposals — queue is empty, should not throw
        await group.ApplyCommittedEntriesAsync(default);

        group.ProcessedIndex.ShouldBe(0);
        group.PendingCommits.ShouldBe(0);
    }

    // ---------------------------------------------------------------
    // Go: raft.go CreateSnapshotCheckpoint — snapshot lifecycle
    // ---------------------------------------------------------------

    [Fact]
    public async Task CheckpointAsync_creates_snapshot_at_current_state()
    {
        var group = new StreamReplicaGroup("SNAP", replicas: 3);

        await group.ProposeAsync("entry.1", default);
        await group.ProposeAsync("entry.2", default);

        var snapshot = await group.CheckpointAsync(default);

        snapshot.ShouldNotBeNull();
        snapshot.LastIncludedIndex.ShouldBeGreaterThan(0);
    }

    [Fact]
    public async Task CheckpointAsync_snapshot_index_matches_applied_index()
    {
        var group = new StreamReplicaGroup("SNAPIDX", replicas: 1);

        await group.ProposeAsync("record.1", default);
        await group.ProposeAsync("record.2", default);

        var snapshot = await group.CheckpointAsync(default);

        snapshot.LastIncludedIndex.ShouldBe(group.Leader.AppliedIndex);
    }

    // ---------------------------------------------------------------
    // Go: raft.go DrainAndReplaySnapshot — restore lifecycle
    // ---------------------------------------------------------------

    [Fact]
    public async Task RestoreFromSnapshotAsync_restores_state()
    {
        var group = new StreamReplicaGroup("RESTORE", replicas: 3);

        await group.ProposeAsync("pre.1", default);
        await group.ProposeAsync("pre.2", default);

        var snapshot = await group.CheckpointAsync(default);

        // Advance state further after snapshot
        await group.ProposeAsync("post.1", default);

        // Restore: should drain queue and roll back to snapshot state
        await group.RestoreFromSnapshotAsync(snapshot, default);

        // After restore the commit index reflects the snapshot
        group.CommitIndex.ShouldBe(snapshot.LastIncludedIndex);
        // Pending commits should be drained
        group.PendingCommits.ShouldBe(0);
    }

    [Fact]
    public async Task RestoreFromSnapshotAsync_drains_pending_commits()
    {
        var group = new StreamReplicaGroup("DRAIN", replicas: 3);

        // Propose several entries so queue has items
        await group.ProposeAsync("queued.1", default);
        await group.ProposeAsync("queued.2", default);
        await group.ProposeAsync("queued.3", default);

        group.PendingCommits.ShouldBeGreaterThan(0);

        var snapshot = new RaftSnapshot
        {
            LastIncludedIndex = 3,
            LastIncludedTerm = group.Leader.Term,
        };

        await group.RestoreFromSnapshotAsync(snapshot, default);

        group.PendingCommits.ShouldBe(0);
    }

    // ---------------------------------------------------------------
    // Go: raft.go:150-160 — PendingCommits reflects commit queue depth
    // ---------------------------------------------------------------

    [Fact]
    public async Task PendingCommits_reflects_commit_queue_depth()
    {
        var group = new StreamReplicaGroup("QUEUE", replicas: 3);

        group.PendingCommits.ShouldBe(0);

        await group.ProposeAsync("q.1", default);
        group.PendingCommits.ShouldBe(1);

        await group.ProposeAsync("q.2", default);
        group.PendingCommits.ShouldBe(2);

        await group.ApplyCommittedEntriesAsync(default);
        group.PendingCommits.ShouldBe(0);
    }

    // ---------------------------------------------------------------
    // Go: raft.go applied/processed tracking — CommitIndex and ProcessedIndex
    // ---------------------------------------------------------------

    [Fact]
    public async Task CommitIndex_and_ProcessedIndex_track_through_the_group()
    {
        var group = new StreamReplicaGroup("INDICES", replicas: 3);

        group.CommitIndex.ShouldBe(0);
        group.ProcessedIndex.ShouldBe(0);

        await group.ProposeAsync("step.1", default);
        group.CommitIndex.ShouldBe(1);
        // Not yet applied
        group.ProcessedIndex.ShouldBe(0);

        await group.ApplyCommittedEntriesAsync(default);
        group.ProcessedIndex.ShouldBe(1);

        await group.ProposeAsync("step.2", default);
        group.CommitIndex.ShouldBe(2);
        group.ProcessedIndex.ShouldBe(1); // still only first entry applied

        await group.ApplyCommittedEntriesAsync(default);
        group.ProcessedIndex.ShouldBe(2);
    }

    [Fact]
    public void CommitIndex_initially_zero_for_fresh_group()
    {
        var group = new StreamReplicaGroup("FRESH", replicas: 5);

        group.CommitIndex.ShouldBe(0);
        group.ProcessedIndex.ShouldBe(0);
        group.PendingCommits.ShouldBe(0);
    }
}