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

using System.IO.Hashing;
using NATS.Server.Raft;

namespace NATS.Server.Tests.Raft;

/// <summary>
/// Tests for Gap 8.3: chunk-based snapshot streaming with CRC32 validation.
///
/// Covers <see cref="SnapshotChunkEnumerator"/> enumeration/CRC behaviour and
/// the <see cref="RaftNode.InstallSnapshotFromChunksAsync"/> CRC validation path.
///
/// Go reference: raft.go:3500-3700 (installSnapshot chunked transfer + CRC validation)
/// </summary>
public class RaftSnapshotStreamingTests
{
    // -----------------------------------------------------------------------
    // SnapshotChunkEnumerator tests
    // -----------------------------------------------------------------------

    [Fact]
    public void SnapshotChunkEnumerator_yields_correct_chunk_count()
    {
        // 200 KB of data at 64 KB per chunk → ceil(200/64) = 4 chunks
        // Go reference: raft.go snapshotChunkSize chunking logic
        const int dataSize = 200 * 1024;
        var data = new byte[dataSize];
        Random.Shared.NextBytes(data);

        var enumerator = new SnapshotChunkEnumerator(data, chunkSize: 65536);
        var chunks = enumerator.ToList();

        chunks.Count.ShouldBe(4);
        enumerator.ChunkCount.ShouldBe(4);
    }

    [Fact]
    public void SnapshotChunkEnumerator_chunks_reassemble_to_original()
    {
        // Concatenating all chunks must reproduce the original byte array exactly
        // Go reference: raft.go installSnapshot chunk reassembly
        const int dataSize = 150 * 1024; // 150 KB → 3 chunks at 64 KB
        var original = new byte[dataSize];
        Random.Shared.NextBytes(original);

        var enumerator = new SnapshotChunkEnumerator(original, chunkSize: 65536);
        var assembled = enumerator.SelectMany(c => c).ToArray();

        assembled.Length.ShouldBe(original.Length);
        assembled.ShouldBe(original);
    }

    [Fact]
    public void SnapshotChunkEnumerator_crc32_matches()
    {
        // The CRC32 reported by the enumerator must equal the CRC32 computed
        // directly over the original data — proving it covers the full payload
        // Go reference: raft.go installSnapshot CRC32 computation
        var data = new byte[100 * 1024];
        Random.Shared.NextBytes(data);

        var enumerator = new SnapshotChunkEnumerator(data, chunkSize: 65536);

        var expectedCrc = new Crc32();
        expectedCrc.Append(data);
        var expected = expectedCrc.GetCurrentHashAsUInt32();

        enumerator.Crc32Value.ShouldBe(expected);
    }

    [Fact]
    public void SnapshotChunkEnumerator_single_chunk_for_small_data()
    {
        // Data that fits in a single chunk — only one chunk should be yielded,
        // and it should be identical to the input
        // Go reference: raft.go installSnapshot — single-chunk case
        var data = new byte[] { 1, 2, 3, 4, 5, 10, 20, 30 };

        var enumerator = new SnapshotChunkEnumerator(data, chunkSize: 65536);
        var chunks = enumerator.ToList();

        chunks.Count.ShouldBe(1);
        enumerator.ChunkCount.ShouldBe(1);
        chunks[0].ShouldBe(data);
    }

    [Fact]
    public void SnapshotChunkEnumerator_last_chunk_is_remainder()
    {
        // 10 bytes with chunk size 3 → chunks of [3, 3, 3, 1]
        var data = Enumerable.Range(0, 10).Select(i => (byte)i).ToArray();

        var enumerator = new SnapshotChunkEnumerator(data, chunkSize: 3);
        var chunks = enumerator.ToList();

        chunks.Count.ShouldBe(4);
        chunks[0].Length.ShouldBe(3);
        chunks[1].Length.ShouldBe(3);
        chunks[2].Length.ShouldBe(3);
        chunks[3].Length.ShouldBe(1); // remainder
        chunks[3][0].ShouldBe((byte)9);
    }

    [Fact]
    public void SnapshotChunkEnumerator_crc32_is_stable_across_multiple_reads()
    {
        // CRC32Value must return the same value on every call (cached)
        var data = new byte[1024];
        Random.Shared.NextBytes(data);

        var enumerator = new SnapshotChunkEnumerator(data, chunkSize: 256);

        var first = enumerator.Crc32Value;
        var second = enumerator.Crc32Value;
        var third = enumerator.Crc32Value;

        second.ShouldBe(first);
        third.ShouldBe(first);
    }

    // -----------------------------------------------------------------------
    // InstallSnapshotFromChunksAsync CRC32 validation tests
    // -----------------------------------------------------------------------

    [Fact]
    public async Task InstallSnapshot_assembles_chunks_into_snapshot()
    {
        // Snapshot assembled from multiple chunks should produce the correct
        // LastIncludedIndex, LastIncludedTerm, and Data on the node.
        // Go reference: raft.go:3500-3700 installSnapshot
        var node = new RaftNode("n1");

        var data = new byte[] { 10, 20, 30, 40, 50 };
        var enumerator = new SnapshotChunkEnumerator(data, chunkSize: 2);
        var chunks = enumerator.ToList();

        await node.InstallSnapshotFromChunksAsync(
            chunks,
            snapshotIndex: 42,
            snapshotTerm: 7,
            ct: default);

        node.AppliedIndex.ShouldBe(42);
        node.CommitIndex.ShouldBe(42);
    }

    [Fact]
    public async Task InstallSnapshot_validates_crc32_success()
    {
        // When the correct CRC32 is supplied the install should complete without error.
        // Go reference: raft.go installSnapshot CRC validation
        var node = new RaftNode("n1");

        var data = new byte[256];
        Random.Shared.NextBytes(data);
        var enumerator = new SnapshotChunkEnumerator(data, chunkSize: 64);
        var crc = enumerator.Crc32Value;
        var chunks = enumerator.ToList();

        // Should not throw
        await node.InstallSnapshotFromChunksAsync(
            chunks,
            snapshotIndex: 10,
            snapshotTerm: 2,
            ct: default,
            expectedCrc32: crc);

        node.AppliedIndex.ShouldBe(10);
    }

    [Fact]
    public async Task InstallSnapshot_validates_crc32_throws_on_mismatch()
    {
        // A wrong CRC32 must cause InvalidDataException before any state is mutated.
        // Go reference: raft.go installSnapshot CRC mismatch → abort
        var node = new RaftNode("n1");
        node.Log.Append(1, "cmd-1"); // pre-existing state

        var data = new byte[] { 1, 2, 3, 4, 5, 6, 7, 8 };
        var chunks = new[] { data }; // single chunk

        const uint wrongCrc = 0xDEADBEEF;

        var ex = await Should.ThrowAsync<InvalidDataException>(async () =>
            await node.InstallSnapshotFromChunksAsync(
                chunks,
                snapshotIndex: 99,
                snapshotTerm: 5,
                ct: default,
                expectedCrc32: wrongCrc));

        ex.Message.ShouldContain("CRC32");
        ex.Message.ShouldContain("DEADBEEF");

        // State must NOT have been mutated since CRC failed before any writes
        node.AppliedIndex.ShouldBe(0);
        node.CommitIndex.ShouldBe(0);
    }

    [Fact]
    public async Task InstallSnapshot_no_crc_parameter_installs_without_validation()
    {
        // When expectedCrc32 is omitted (null), no validation occurs and any data installs.
        // Go reference: raft.go optional CRC field (backward compat)
        var node = new RaftNode("n1");
        var chunks = new[] { new byte[] { 7, 8, 9 } };

        // Should not throw even with no CRC supplied
        await node.InstallSnapshotFromChunksAsync(
            chunks,
            snapshotIndex: 5,
            snapshotTerm: 1,
            ct: default);

        node.AppliedIndex.ShouldBe(5);
        node.CommitIndex.ShouldBe(5);
    }

    // -----------------------------------------------------------------------
    // RaftInstallSnapshotChunkWire encode/decode roundtrip tests
    // -----------------------------------------------------------------------

    [Fact]
    public void SnapshotChunkWire_roundtrip_with_data()
    {
        // Encode and decode a chunk message and verify all fields survive the roundtrip.
        // Go reference: raft.go wire format for InstallSnapshot RPC chunks
        var payload = new byte[] { 0xAA, 0xBB, 0xCC, 0xDD };
        var wire = new RaftInstallSnapshotChunkWire(
            SnapshotIndex: 42UL,
            SnapshotTerm: 3U,
            ChunkIndex: 1U,
            TotalChunks: 4U,
            Crc32: 0x12345678U,
            Data: payload);

        var encoded = wire.Encode();
        var decoded = RaftInstallSnapshotChunkWire.Decode(encoded);

        decoded.SnapshotIndex.ShouldBe(42UL);
        decoded.SnapshotTerm.ShouldBe(3U);
        decoded.ChunkIndex.ShouldBe(1U);
        decoded.TotalChunks.ShouldBe(4U);
        decoded.Crc32.ShouldBe(0x12345678U);
        decoded.Data.ShouldBe(payload);
    }

    [Fact]
    public void SnapshotChunkWire_header_length_is_24_bytes()
    {
        // Header must be exactly 24 bytes as documented in the wire format.
        RaftInstallSnapshotChunkWire.HeaderLen.ShouldBe(24);
    }

    [Fact]
    public void SnapshotChunkWire_encode_total_length_is_header_plus_data()
    {
        var data = new byte[100];
        var wire = new RaftInstallSnapshotChunkWire(1UL, 1U, 0U, 1U, 0U, data);
        wire.Encode().Length.ShouldBe(RaftInstallSnapshotChunkWire.HeaderLen + 100);
    }

    [Fact]
    public void SnapshotChunkWire_decode_throws_on_short_buffer()
    {
        // Buffers shorter than the header should throw ArgumentException
        var tooShort = new byte[10]; // < 24 bytes
        Should.Throw<ArgumentException>(() => RaftInstallSnapshotChunkWire.Decode(tooShort));
    }

    [Fact]
    public void SnapshotChunkWire_roundtrip_empty_payload()
    {
        // A header-only message (no chunk data) should encode and decode cleanly.
        var wire = new RaftInstallSnapshotChunkWire(
            SnapshotIndex: 0UL,
            SnapshotTerm: 0U,
            ChunkIndex: 0U,
            TotalChunks: 0U,
            Crc32: 0U,
            Data: []);

        var encoded = wire.Encode();
        encoded.Length.ShouldBe(RaftInstallSnapshotChunkWire.HeaderLen);

        var decoded = RaftInstallSnapshotChunkWire.Decode(encoded);
        decoded.Data.Length.ShouldBe(0);
    }
}