natsdotnet/tests/NATS.Server.JetStream.Tests/JetStream/Snapshots/StreamSnapshotTests.cs

// Go reference: server/jetstream_api.go — jsStreamSnapshotT / jsStreamRestoreT
// Tests for TAR-based stream snapshot/restore with Snappy (S2) compression.

using System.Formats.Tar;
using System.Text;
using System.Text.Json;
using IronSnappy;
using NATS.Server.JetStream;
using NATS.Server.JetStream.Models;
using NATS.Server.JetStream.Snapshots;
using NATS.Server.JetStream.Storage;
using Shouldly;

namespace NATS.Server.JetStream.Tests.JetStream.Snapshots;

public sealed class StreamSnapshotTests
{
    // ──────────────────────────────────────────────────────────────────────────
    // Helpers
    // ──────────────────────────────────────────────────────────────────────────

    private static StreamHandle MakeHandle(string name = "TEST")
    {
        var config = new StreamConfig { Name = name };
        IStreamStore store = new MemStore(config);
        return new StreamHandle(config, store);
    }

    private static async Task StoreAsync(StreamHandle h, string subject, string payload)
        => await h.Store.AppendAsync(subject, Encoding.UTF8.GetBytes(payload), default);

    private static async Task<(byte[] tarBytes, List<string> entryNames)> DecompressAndListEntries(
        byte[] snappyBytes)
    {
        var tarBytes = Snappy.Decode(snappyBytes);
        using var ms     = new MemoryStream(tarBytes);
        using var reader = new TarReader(ms, leaveOpen: false);

        var names = new List<string>();
        TarEntry? entry;
        while ((entry = await reader.GetNextEntryAsync()) is not null)
            names.Add(entry.Name);

        return (tarBytes, names);
    }

    // ──────────────────────────────────────────────────────────────────────────
    // Test 1 — stream.json is present in the TAR
    // ──────────────────────────────────────────────────────────────────────────

    // Go ref: server/jetstream_api.go:jsStreamSnapshotT — snapshot includes config JSON.
    [Fact]
    public async Task CreateTarSnapshot_includes_stream_config()
    {
        var svc    = new StreamSnapshotService();
        var handle = MakeHandle("MYSTREAM");

        var compressed = await svc.CreateTarSnapshotAsync(handle, default);

        var (_, names) = await DecompressAndListEntries(compressed);
        names.ShouldContain("stream.json");
    }

    // ──────────────────────────────────────────────────────────────────────────
    // Test 2 — all messages appear in the TAR
    // ──────────────────────────────────────────────────────────────────────────

    // Go ref: server/jetstream_api.go:jsStreamSnapshotT — snapshot contains all stored messages.
    [Fact]
    public async Task CreateTarSnapshot_includes_all_messages()
    {
        var svc    = new StreamSnapshotService();
        var handle = MakeHandle();

        for (var i = 1; i <= 5; i++)
            await StoreAsync(handle, $"foo.{i}", $"payload-{i}");

        var compressed = await svc.CreateTarSnapshotAsync(handle, default);

        var (_, names) = await DecompressAndListEntries(compressed);
        var messageEntries = names.Where(n => n.StartsWith("messages/")).ToList();
        messageEntries.Count.ShouldBe(5);
    }

    // ──────────────────────────────────────────────────────────────────────────
    // Test 3 — the snapshot bytes are Snappy-compressed
    // ──────────────────────────────────────────────────────────────────────────

    // Go ref: server/filestore.go — S2 / Snappy compression is used for snapshots.
    [Fact]
    public async Task CreateTarSnapshot_compresses_with_snappy()
    {
        var svc    = new StreamSnapshotService();
        var handle = MakeHandle();
        await StoreAsync(handle, "test.subject", "test payload data");

        var compressed = await svc.CreateTarSnapshotAsync(handle, default);

        // Round-trip decode must succeed — this verifies it is valid Snappy data.
        var decoded = Should.NotThrow(() => Snappy.Decode(compressed));
        decoded.ShouldNotBeEmpty();
    }

    // ──────────────────────────────────────────────────────────────────────────
    // Test 4 — restore replays messages back into the store
    // ──────────────────────────────────────────────────────────────────────────

    // Go ref: server/jetstream_api.go:jsStreamRestoreT — restore re-populates the store.
    [Fact]
    public async Task RestoreTarSnapshot_replays_messages()
    {
        var svc    = new StreamSnapshotService();
        var handle = MakeHandle();

        await StoreAsync(handle, "a.1", "first");
        await StoreAsync(handle, "a.2", "second");
        await StoreAsync(handle, "a.3", "third");

        var compressed = await svc.CreateTarSnapshotAsync(handle, default);

        // Purge the store to simulate data loss.
        await handle.Store.PurgeAsync(default);
        var stateAfterPurge = await handle.Store.GetStateAsync(default);
        stateAfterPurge.Messages.ShouldBe(0UL);

        // Restore from the snapshot.
        await svc.RestoreTarSnapshotAsync(handle, compressed, default);

        var stateAfterRestore = await handle.Store.GetStateAsync(default);
        stateAfterRestore.Messages.ShouldBe(3UL);
    }

    // ──────────────────────────────────────────────────────────────────────────
    // Test 5 — restore returns correct stats
    // ──────────────────────────────────────────────────────────────────────────

    // Go ref: server/jetstream_api.go:jsStreamRestoreT — response reports bytes/messages restored.
    [Fact]
    public async Task RestoreTarSnapshot_returns_correct_stats()
    {
        var svc    = new StreamSnapshotService();
        var handle = MakeHandle("STATSSTREAM");

        await StoreAsync(handle, "s.1", "hello");
        await StoreAsync(handle, "s.2", "world");

        var compressed = await svc.CreateTarSnapshotAsync(handle, default);
        await handle.Store.PurgeAsync(default);

        var result = await svc.RestoreTarSnapshotAsync(handle, compressed, default);

        result.StreamName.ShouldBe("STATSSTREAM");
        result.MessagesRestored.ShouldBe(2);
        result.BytesRestored.ShouldBeGreaterThan(0L);
    }

    // ──────────────────────────────────────────────────────────────────────────
    // Test 6 — restore rejects a TAR that is missing stream.json
    // ──────────────────────────────────────────────────────────────────────────

    // Go ref: server/jetstream_api.go:jsStreamRestoreT — invalid snapshot returns an error.
    [Fact]
    public async Task RestoreTarSnapshot_validates_stream_json()
    {
        var svc    = new StreamSnapshotService();
        var handle = MakeHandle();

        // Build a TAR that contains no stream.json entry.
        using var tarBuffer = new MemoryStream();
        using (var writer = new TarWriter(tarBuffer, TarEntryFormat.Pax, leaveOpen: true))
        {
            var bogus = new PaxTarEntry(TarEntryType.RegularFile, "other/file.txt")
            {
                DataStream = new MemoryStream("not a config"u8.ToArray()),
            };
            await writer.WriteEntryAsync(bogus);
        }

        var badSnappy = Snappy.Encode(tarBuffer.ToArray());

        await Should.ThrowAsync<InvalidOperationException>(
            () => svc.RestoreTarSnapshotAsync(handle, badSnappy, default));
    }

    // ──────────────────────────────────────────────────────────────────────────
    // Test 7 — snapshot with deadline completes when deadline is generous
    // ──────────────────────────────────────────────────────────────────────────

    // Go ref: server/filestore.go — snapshot creation must respect deadlines.
    [Fact]
    public async Task CreateTarSnapshotWithDeadline_completes_within_deadline()
    {
        var svc    = new StreamSnapshotService();
        var handle = MakeHandle();
        await StoreAsync(handle, "d.1", "data");

        var compressed = await svc.CreateTarSnapshotWithDeadlineAsync(
            handle,
            TimeSpan.FromSeconds(30),
            default);

        compressed.ShouldNotBeEmpty();
    }

    // ──────────────────────────────────────────────────────────────────────────
    // Test 8 — snapshot of an empty stream produces a valid archive
    // ──────────────────────────────────────────────────────────────────────────

    // Go ref: server/jetstream_api.go:jsStreamSnapshotT — empty stream is a valid snapshot.
    [Fact]
    public async Task CreateTarSnapshot_empty_stream()
    {
        var svc    = new StreamSnapshotService();
        var handle = MakeHandle("EMPTY");

        var compressed = await svc.CreateTarSnapshotAsync(handle, default);

        compressed.ShouldNotBeEmpty();
        var (_, names) = await DecompressAndListEntries(compressed);
        names.ShouldContain("stream.json");
        names.Where(n => n.StartsWith("messages/")).ShouldBeEmpty();
    }

    // ──────────────────────────────────────────────────────────────────────────
    // Test 9 — round-trip preserves message subjects
    // ──────────────────────────────────────────────────────────────────────────

    // Go ref: server/jetstream_api.go — snapshot/restore is lossless for subjects.
    [Fact]
    public async Task Roundtrip_snapshot_preserves_message_subjects()
    {
        var svc    = new StreamSnapshotService();
        var handle = MakeHandle("SUBJECTS");

        var subjects = new[] { "alpha.1", "beta.2", "gamma.3" };
        foreach (var s in subjects)
            await StoreAsync(handle, s, $"payload for {s}");

        var compressed = await svc.CreateTarSnapshotAsync(handle, default);

        // Restore to a fresh store.
        var newConfig = new StreamConfig { Name = "SUBJECTS" };
        IStreamStore newStore = new MemStore(newConfig);
        var newHandle = new StreamHandle(newConfig, newStore);

        await svc.RestoreTarSnapshotAsync(newHandle, compressed, default);

        var restored = await newStore.ListAsync(default);
        var restoredSubjects = restored.Select(m => m.Subject).OrderBy(x => x).ToArray();
        var expectedSubjects = subjects.OrderBy(x => x).ToArray();
        restoredSubjects.ShouldBe(expectedSubjects);
    }

    // ──────────────────────────────────────────────────────────────────────────
    // Test 10 — round-trip preserves message payloads
    // ──────────────────────────────────────────────────────────────────────────

    // Go ref: server/jetstream_api.go — snapshot/restore is lossless for payload bytes.
    [Fact]
    public async Task Roundtrip_snapshot_preserves_message_payloads()
    {
        var svc    = new StreamSnapshotService();
        var handle = MakeHandle("PAYLOADS");

        var payloads = new[] { "hello world"u8.ToArray(), new byte[] { 0, 1, 2, 3, 255 } };
        await handle.Store.AppendAsync("p.1", payloads[0], default);
        await handle.Store.AppendAsync("p.2", payloads[1], default);

        var compressed = await svc.CreateTarSnapshotAsync(handle, default);

        // Restore to a fresh store.
        var newConfig = new StreamConfig { Name = "PAYLOADS" };
        IStreamStore newStore = new MemStore(newConfig);
        var newHandle = new StreamHandle(newConfig, newStore);

        await svc.RestoreTarSnapshotAsync(newHandle, compressed, default);

        var restored = (await newStore.ListAsync(default)).OrderBy(m => m.Sequence).ToList();
        restored.Count.ShouldBe(2);
        restored[0].Payload.ToArray().ShouldBe(payloads[0]);
        restored[1].Payload.ToArray().ShouldBe(payloads[1]);
    }
}