natsdotnet/tests/NATS.Server.Raft.Tests/Raft/RaftSnapshotCheckpointTests.cs

using NATS.Server.Raft;

namespace NATS.Server.Raft.Tests.Raft;

/// <summary>
/// Tests for B5: Snapshot Checkpoints and Log Compaction.
/// Go reference: raft.go:3200-3400 (CreateSnapshotCheckpoint), raft.go:3500-3700 (installSnapshot).
/// </summary>
public class RaftSnapshotCheckpointTests
{
    // -- Helpers --

    private static (RaftNode[] nodes, InMemoryRaftTransport transport) CreateCluster(int size)
    {
        var transport = new InMemoryRaftTransport();
        var nodes = Enumerable.Range(1, size)
            .Select(i => new RaftNode($"n{i}", transport))
            .ToArray();
        foreach (var node in nodes)
        {
            transport.Register(node);
            node.ConfigureCluster(nodes);
        }

        return (nodes, transport);
    }

    private static RaftNode ElectLeader(RaftNode[] nodes)
    {
        var candidate = nodes[0];
        candidate.StartElection(nodes.Length);
        foreach (var voter in nodes.Skip(1))
            candidate.ReceiveVote(voter.GrantVote(candidate.Term, candidate.Id), nodes.Length);
        return candidate;
    }

    // -- RaftSnapshotCheckpoint type tests --

    [Fact]
    public void Checkpoint_creation_with_data()
    {
        var checkpoint = new RaftSnapshotCheckpoint
        {
            SnapshotIndex = 10,
            SnapshotTerm = 2,
            Data = [1, 2, 3, 4, 5],
        };

        checkpoint.SnapshotIndex.ShouldBe(10);
        checkpoint.SnapshotTerm.ShouldBe(2);
        checkpoint.Data.Length.ShouldBe(5);
        checkpoint.IsComplete.ShouldBeFalse();
    }

    [Fact]
    public void Chunk_assembly_single_chunk()
    {
        var checkpoint = new RaftSnapshotCheckpoint
        {
            SnapshotIndex = 5,
            SnapshotTerm = 1,
        };

        checkpoint.AddChunk([10, 20, 30]);
        var result = checkpoint.Assemble();

        result.Length.ShouldBe(3);
        result[0].ShouldBe((byte)10);
        result[1].ShouldBe((byte)20);
        result[2].ShouldBe((byte)30);
        checkpoint.IsComplete.ShouldBeTrue();
    }

    [Fact]
    public void Chunk_assembly_multiple_chunks()
    {
        var checkpoint = new RaftSnapshotCheckpoint
        {
            SnapshotIndex = 5,
            SnapshotTerm = 1,
        };

        checkpoint.AddChunk([1, 2]);
        checkpoint.AddChunk([3, 4, 5]);
        checkpoint.AddChunk([6]);

        var result = checkpoint.Assemble();
        result.Length.ShouldBe(6);
        result.ShouldBe(new byte[] { 1, 2, 3, 4, 5, 6 });
        checkpoint.IsComplete.ShouldBeTrue();
    }

    [Fact]
    public void Chunk_assembly_empty_returns_data()
    {
        // When no chunks added, Assemble returns the initial Data property
        var checkpoint = new RaftSnapshotCheckpoint
        {
            SnapshotIndex = 5,
            SnapshotTerm = 1,
            Data = [99, 100],
        };

        var result = checkpoint.Assemble();
        result.ShouldBe(new byte[] { 99, 100 });
        checkpoint.IsComplete.ShouldBeFalse(); // no chunks to assemble
    }

    // -- RaftLog.Compact tests --

    [Fact]
    public void CompactLog_removes_old_entries()
    {
        // Go reference: raft.go WAL compact
        var log = new RaftLog();
        log.Append(1, "cmd-1");
        log.Append(1, "cmd-2");
        log.Append(1, "cmd-3");
        log.Append(2, "cmd-4");
        log.Entries.Count.ShouldBe(4);

        // Compact up to index 2 — entries 1 and 2 should be removed
        log.Compact(2);
        log.Entries.Count.ShouldBe(2);
        log.Entries[0].Index.ShouldBe(3);
        log.Entries[1].Index.ShouldBe(4);
    }

    [Fact]
    public void CompactLog_updates_base_index()
    {
        var log = new RaftLog();
        log.Append(1, "cmd-1");
        log.Append(1, "cmd-2");
        log.Append(1, "cmd-3");

        log.BaseIndex.ShouldBe(0);
        log.Compact(2);
        log.BaseIndex.ShouldBe(2);
    }

    [Fact]
    public void CompactLog_with_no_entries_is_noop()
    {
        var log = new RaftLog();
        log.Entries.Count.ShouldBe(0);
        log.BaseIndex.ShouldBe(0);

        // Should not throw or change anything
        log.Compact(5);
        log.Entries.Count.ShouldBe(0);
        log.BaseIndex.ShouldBe(0);
    }

    [Fact]
    public void CompactLog_preserves_append_indexing()
    {
        // After compaction, new appends should continue from the correct index
        var log = new RaftLog();
        log.Append(1, "cmd-1");
        log.Append(1, "cmd-2");
        log.Append(1, "cmd-3");

        log.Compact(2);
        log.BaseIndex.ShouldBe(2);

        // New entry should get index 4 (baseIndex 2 + 1 remaining entry + 1)
        var newEntry = log.Append(2, "cmd-4");
        newEntry.Index.ShouldBe(4);
    }

    // -- Streaming snapshot install on RaftNode --

    [Fact]
    public async Task Streaming_snapshot_install_from_chunks()
    {
        // Go reference: raft.go:3500-3700 (installSnapshot with chunked transfer)
        var node = new RaftNode("n1");
        node.Log.Append(1, "cmd-1");
        node.Log.Append(1, "cmd-2");
        node.Log.Append(1, "cmd-3");

        byte[][] chunks = [[1, 2, 3], [4, 5, 6]];
        await node.InstallSnapshotFromChunksAsync(chunks, snapshotIndex: 10, snapshotTerm: 3, default);

        // Log should be replaced (entries cleared, base index set to snapshot)
        node.Log.Entries.Count.ShouldBe(0);
        node.AppliedIndex.ShouldBe(10);
        node.CommitIndex.ShouldBe(10);
    }

    [Fact]
    public async Task Log_after_compaction_starts_at_correct_index()
    {
        // After snapshot + compaction, new entries should continue from the right index
        var (nodes, _) = CreateCluster(3);
        var leader = ElectLeader(nodes);

        await leader.ProposeAsync("cmd-1", default);
        await leader.ProposeAsync("cmd-2", default);
        await leader.ProposeAsync("cmd-3", default);

        leader.Log.Entries.Count.ShouldBe(3);

        // Create snapshot at current applied index and compact
        var snapshot = await leader.CreateSnapshotCheckpointAsync(default);
        snapshot.LastIncludedIndex.ShouldBe(leader.AppliedIndex);

        // Log should now be empty (all entries covered by snapshot)
        leader.Log.Entries.Count.ShouldBe(0);
        leader.Log.BaseIndex.ShouldBe(leader.AppliedIndex);

        // New entries should continue from the right index
        var index4 = await leader.ProposeAsync("cmd-4", default);
        index4.ShouldBe(leader.AppliedIndex); // should be appliedIndex after new propose
        leader.Log.Entries.Count.ShouldBe(1);
    }

    // -- CompactLogAsync on RaftNode --

    [Fact]
    public async Task CompactLogAsync_compacts_up_to_applied_index()
    {
        var (nodes, _) = CreateCluster(3);
        var leader = ElectLeader(nodes);

        await leader.ProposeAsync("cmd-1", default);
        await leader.ProposeAsync("cmd-2", default);
        await leader.ProposeAsync("cmd-3", default);

        leader.Log.Entries.Count.ShouldBe(3);
        var appliedIndex = leader.AppliedIndex;
        appliedIndex.ShouldBeGreaterThan(0);

        await leader.CompactLogAsync(default);

        // All entries up to applied index should be compacted
        leader.Log.BaseIndex.ShouldBe(appliedIndex);
        leader.Log.Entries.Count.ShouldBe(0);
    }

    [Fact]
    public async Task CompactLogAsync_noop_when_nothing_applied()
    {
        var node = new RaftNode("n1");
        node.AppliedIndex.ShouldBe(0);

        // Should be a no-op — nothing to compact
        await node.CompactLogAsync(default);
        node.Log.BaseIndex.ShouldBe(0);
        node.Log.Entries.Count.ShouldBe(0);
    }
}