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

using System.Text.Json;
using NATS.Server.Raft;

namespace NATS.Server.Tests.Raft;

/// <summary>
/// Snapshot tests covering creation, restore, transfer, membership changes during
/// snapshot, snapshot store persistence, and leader/follower catchup via snapshots.
/// Go: TestNRGSnapshotAndRestart, TestNRGRemoveLeaderPeerDeadlockBug,
///     TestNRGLeaderTransfer in server/raft_test.go.
/// </summary>
public class RaftSnapshotTests
{
    // -- Helpers (self-contained) --

    private static (RaftNode leader, RaftNode[] followers) CreateLeaderWithFollowers(int followerCount)
    {
        var total = followerCount + 1;
        var nodes = Enumerable.Range(1, total)
            .Select(i => new RaftNode($"n{i}"))
            .ToArray();
        foreach (var node in nodes)
            node.ConfigureCluster(nodes);

        var candidate = nodes[0];
        candidate.StartElection(total);
        foreach (var voter in nodes.Skip(1))
            candidate.ReceiveVote(voter.GrantVote(candidate.Term, candidate.Id), total);

        return (candidate, nodes.Skip(1).ToArray());
    }

    private static (RaftNode leader, RaftNode[] followers, InMemoryRaftTransport transport) CreateTransportCluster(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);
        }

        var candidate = nodes[0];
        candidate.StartElection(size);
        foreach (var voter in nodes.Skip(1))
            candidate.ReceiveVote(voter.GrantVote(candidate.Term, candidate.Id), size);

        return (candidate, nodes.Skip(1).ToArray(), transport);
    }

    // Go: TestNRGSnapshotAndRestart server/raft_test.go:49 — snapshot creation
    [Fact]
    public async Task Create_snapshot_captures_applied_index_and_term()
    {
        var (leader, _) = CreateLeaderWithFollowers(2);
        await leader.ProposeAsync("cmd-1", default);
        await leader.ProposeAsync("cmd-2", default);

        var snapshot = await leader.CreateSnapshotAsync(default);

        snapshot.LastIncludedIndex.ShouldBe(leader.AppliedIndex);
        snapshot.LastIncludedTerm.ShouldBe(leader.Term);
    }

    // Go: TestNRGSnapshotAndRestart server/raft_test.go:49 — install snapshot
    [Fact]
    public async Task Install_snapshot_updates_applied_index()
    {
        var (leader, followers) = CreateLeaderWithFollowers(2);
        await leader.ProposeAsync("snap-cmd-1", default);
        await leader.ProposeAsync("snap-cmd-2", default);
        await leader.ProposeAsync("snap-cmd-3", default);

        var snapshot = await leader.CreateSnapshotAsync(default);
        var newFollower = new RaftNode("new-follower");

        await newFollower.InstallSnapshotAsync(snapshot, default);

        newFollower.AppliedIndex.ShouldBe(snapshot.LastIncludedIndex);
    }

    // Go: TestNRGSnapshotAndRestart server/raft_test.go:49 — snapshot clears log
    [Fact]
    public async Task Install_snapshot_clears_existing_log()
    {
        var node = new RaftNode("n1");
        node.Log.Append(term: 1, command: "old-1");
        node.Log.Append(term: 1, command: "old-2");
        node.Log.Entries.Count.ShouldBe(2);

        var snapshot = new RaftSnapshot
        {
            LastIncludedIndex = 10,
            LastIncludedTerm = 3,
        };

        await node.InstallSnapshotAsync(snapshot, default);

        node.Log.Entries.Count.ShouldBe(0);
        node.AppliedIndex.ShouldBe(10);
    }

    // Go: TestNRGSnapshotAndRestart server/raft_test.go:49 — new entries after snapshot
    [Fact]
    public async Task Entries_after_snapshot_start_at_correct_index()
    {
        var node = new RaftNode("n1");

        var snapshot = new RaftSnapshot
        {
            LastIncludedIndex = 50,
            LastIncludedTerm = 5,
        };

        await node.InstallSnapshotAsync(snapshot, default);

        var entry = node.Log.Append(term: 6, command: "post-snap");
        entry.Index.ShouldBe(51);
    }

    // Go: TestNRGSnapshotAndRestart server/raft_test.go:49 — snapshot transfer
    [Fact]
    public async Task Snapshot_transfer_via_transport()
    {
        var (leader, followers, transport) = CreateTransportCluster(3);
        await leader.ProposeAsync("entry-1", default);
        await leader.ProposeAsync("entry-2", default);

        var snapshot = await leader.CreateSnapshotAsync(default);

        // Transfer to a follower
        var follower = followers[0];
        await transport.InstallSnapshotAsync(leader.Id, follower.Id, snapshot, default);

        follower.AppliedIndex.ShouldBe(snapshot.LastIncludedIndex);
    }

    // Go: TestNRGSnapshotAndRestart server/raft_test.go:49 — lagging follower catchup
    [Fact]
    public async Task Lagging_follower_catches_up_via_snapshot()
    {
        var (leader, followers) = CreateLeaderWithFollowers(2);

        // Leader has entries, follower is behind
        await leader.ProposeAsync("catchup-1", default);
        await leader.ProposeAsync("catchup-2", default);
        await leader.ProposeAsync("catchup-3", default);

        var laggingFollower = new RaftNode("lagging");
        laggingFollower.AppliedIndex.ShouldBe(0);

        var snapshot = await leader.CreateSnapshotAsync(default);
        await laggingFollower.InstallSnapshotAsync(snapshot, default);

        laggingFollower.AppliedIndex.ShouldBe(leader.AppliedIndex);
    }

    // Go: RaftSnapshotStore — in-memory save/load
    [Fact]
    public async Task Snapshot_store_in_memory_save_and_load()
    {
        var store = new RaftSnapshotStore();

        var snapshot = new RaftSnapshot
        {
            LastIncludedIndex = 42,
            LastIncludedTerm = 7,
            Data = [1, 2, 3],
        };

        await store.SaveAsync(snapshot, default);
        var loaded = await store.LoadAsync(default);

        loaded.ShouldNotBeNull();
        loaded.LastIncludedIndex.ShouldBe(42);
        loaded.LastIncludedTerm.ShouldBe(7);
        loaded.Data.ShouldBe(new byte[] { 1, 2, 3 });
    }

    // Go: RaftSnapshotStore — file-based save/load
    [Fact]
    public async Task Snapshot_store_file_based_persistence()
    {
        var file = Path.Combine(Path.GetTempPath(), $"nats-raft-snap-{Guid.NewGuid():N}.json");

        try
        {
            var store1 = new RaftSnapshotStore(file);
            await store1.SaveAsync(new RaftSnapshot
            {
                LastIncludedIndex = 100,
                LastIncludedTerm = 10,
                Data = [99, 88, 77],
            }, default);

            // New store instance, load from file
            var store2 = new RaftSnapshotStore(file);
            var loaded = await store2.LoadAsync(default);

            loaded.ShouldNotBeNull();
            loaded.LastIncludedIndex.ShouldBe(100);
            loaded.LastIncludedTerm.ShouldBe(10);
            loaded.Data.ShouldBe(new byte[] { 99, 88, 77 });
        }
        finally
        {
            if (File.Exists(file))
                File.Delete(file);
        }
    }

    // Go: RaftSnapshotStore — load from nonexistent returns null
    [Fact]
    public async Task Snapshot_store_load_nonexistent_returns_null()
    {
        var store = new RaftSnapshotStore();
        var loaded = await store.LoadAsync(default);
        loaded.ShouldBeNull();
    }

    // Go: TestNRGRemoveLeaderPeerDeadlockBug server/raft_test.go:1040 — membership add
    [Fact]
    public void Membership_add_member()
    {
        var node = new RaftNode("n1");
        node.Members.ShouldContain("n1"); // self is auto-added

        node.AddMember("n2");
        node.AddMember("n3");
        node.Members.ShouldContain("n2");
        node.Members.ShouldContain("n3");
        node.Members.Count.ShouldBe(3);
    }

    // Go: TestNRGRemoveLeaderPeerDeadlockBug server/raft_test.go:1040 — membership remove
    [Fact]
    public void Membership_remove_member()
    {
        var node = new RaftNode("n1");
        node.AddMember("n2");
        node.AddMember("n3");

        node.RemoveMember("n2");
        node.Members.ShouldNotContain("n2");
        node.Members.ShouldContain("n1");
        node.Members.ShouldContain("n3");
    }

    // Go: TestNRGRemoveLeaderPeerDeadlockBug server/raft_test.go:1040
    [Fact]
    public void Remove_nonexistent_member_is_noop()
    {
        var node = new RaftNode("n1");
        node.RemoveMember("nonexistent"); // should not throw
        node.Members.Count.ShouldBe(1); // still just self
    }

    // Go: ConfigureCluster in RaftNode
    [Fact]
    public void Configure_cluster_sets_members()
    {
        var n1 = new RaftNode("n1");
        var n2 = new RaftNode("n2");
        var n3 = new RaftNode("n3");
        var nodes = new[] { n1, n2, n3 };

        n1.ConfigureCluster(nodes);

        n1.Members.ShouldContain("n1");
        n1.Members.ShouldContain("n2");
        n1.Members.ShouldContain("n3");
    }

    // Go: TestNRGLeaderTransfer server/raft_test.go:377 — leadership transfer
    [Fact]
    public async Task Leadership_transfer_via_stepdown_and_reelection()
    {
        var (leader, followers) = CreateLeaderWithFollowers(2);
        leader.IsLeader.ShouldBeTrue();

        var preferredNode = followers[0];

        // Leader steps down
        leader.RequestStepDown();
        leader.IsLeader.ShouldBeFalse();

        // Preferred node runs election
        var allNodes = new[] { leader }.Concat(followers).ToArray();
        preferredNode.StartElection(allNodes.Length);

        foreach (var voter in allNodes.Where(n => n.Id != preferredNode.Id))
        {
            var vote = voter.GrantVote(preferredNode.Term, preferredNode.Id);
            preferredNode.ReceiveVote(vote, allNodes.Length);
        }

        preferredNode.IsLeader.ShouldBeTrue();
    }

    // Go: TestNRGSnapshotAndRestart server/raft_test.go:49 — snapshot with data payload
    [Fact]
    public void Snapshot_with_large_data_payload()
    {
        var data = new byte[1024 * 64]; // 64KB
        Random.Shared.NextBytes(data);

        var snapshot = new RaftSnapshot
        {
            LastIncludedIndex = 500,
            LastIncludedTerm = 20,
            Data = data,
        };

        snapshot.Data.Length.ShouldBe(1024 * 64);
        snapshot.LastIncludedIndex.ShouldBe(500);
    }

    // Go: TestNRGSnapshotAndRestart server/raft_test.go:49 — snapshot JSON round-trip
    [Fact]
    public void Snapshot_json_serialization_round_trip()
    {
        var data = new byte[] { 10, 20, 30, 40, 50 };
        var snapshot = new RaftSnapshot
        {
            LastIncludedIndex = 75,
            LastIncludedTerm = 8,
            Data = data,
        };

        var json = JsonSerializer.Serialize(snapshot);
        var decoded = JsonSerializer.Deserialize<RaftSnapshot>(json);

        decoded.ShouldNotBeNull();
        decoded.LastIncludedIndex.ShouldBe(75);
        decoded.LastIncludedTerm.ShouldBe(8);
        decoded.Data.ShouldBe(data);
    }

    // Go: TestNRGSnapshotAndRestart server/raft_test.go:49 — full cluster snapshot + restart
    [Fact]
    public async Task Full_cluster_snapshot_and_follower_restart()
    {
        var (leader, followers) = CreateLeaderWithFollowers(2);

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

        var snapshot = await leader.CreateSnapshotAsync(default);

        // Simulate follower restart by installing snapshot on fresh node
        var restartedFollower = new RaftNode("restarted");
        await restartedFollower.InstallSnapshotAsync(snapshot, default);

        restartedFollower.AppliedIndex.ShouldBe(snapshot.LastIncludedIndex);
        restartedFollower.Log.Entries.Count.ShouldBe(0); // log was replaced by snapshot
    }

    // Go: TestNRGSnapshotAndRestart server/raft_test.go:49 — snapshot replaces stale log
    [Fact]
    public async Task Snapshot_replaces_stale_log_entries()
    {
        var node = new RaftNode("n1");
        node.Log.Append(term: 1, command: "stale-1");
        node.Log.Append(term: 1, command: "stale-2");
        node.Log.Append(term: 1, command: "stale-3");

        var snapshot = new RaftSnapshot
        {
            LastIncludedIndex = 100,
            LastIncludedTerm = 5,
        };

        await node.InstallSnapshotAsync(snapshot, default);

        node.Log.Entries.Count.ShouldBe(0);
        node.AppliedIndex.ShouldBe(100);

        // New entries continue from snapshot base
        var newEntry = node.Log.Append(term: 6, command: "fresh");
        newEntry.Index.ShouldBe(101);
    }

    // Go: TestNRGSnapshotAndRestart server/raft_test.go:49 — snapshot store overwrites previous
    [Fact]
    public async Task Snapshot_store_overwrites_previous_snapshot()
    {
        var store = new RaftSnapshotStore();

        await store.SaveAsync(new RaftSnapshot { LastIncludedIndex = 10, LastIncludedTerm = 1 }, default);
        await store.SaveAsync(new RaftSnapshot { LastIncludedIndex = 50, LastIncludedTerm = 3 }, default);

        var loaded = await store.LoadAsync(default);
        loaded.ShouldNotBeNull();
        loaded.LastIncludedIndex.ShouldBe(50);
        loaded.LastIncludedTerm.ShouldBe(3);
    }

    // Go: TestNRGSnapshotAndRestart server/raft_test.go:49 — node state after multiple snapshots
    [Fact]
    public async Task Multiple_snapshot_installs_advance_applied_index()
    {
        var node = new RaftNode("n1");

        await node.InstallSnapshotAsync(new RaftSnapshot
        {
            LastIncludedIndex = 10,
            LastIncludedTerm = 1,
        }, default);
        node.AppliedIndex.ShouldBe(10);

        await node.InstallSnapshotAsync(new RaftSnapshot
        {
            LastIncludedIndex = 50,
            LastIncludedTerm = 3,
        }, default);
        node.AppliedIndex.ShouldBe(50);

        // Entries start after latest snapshot
        var entry = node.Log.Append(term: 4, command: "after-second-snap");
        entry.Index.ShouldBe(51);
    }
}