natsdotnet/tests/NATS.Server.JetStream.Tests/JetStream/Storage/MemStoreTests.cs

// Reference: golang/nats-server/server/memstore_test.go and filestore_test.go
// Tests ported from: TestMemStoreBasics, TestMemStorePurge, TestMemStoreMsgHeaders,
//                    TestMemStoreTimeStamps, TestMemStoreEraseMsg,
//                    TestMemStoreMsgLimit, TestMemStoreBytesLimit,
//                    TestMemStoreAgeLimit, plus parity tests matching
//                    filestore behavior in MemStore.

using System.Text;
using NATS.Server.JetStream.Storage;

namespace NATS.Server.JetStream.Tests.JetStream.Storage;

public sealed class MemStoreTests
{
    // Go: TestMemStoreBasics server/memstore_test.go
    [Fact]
    public async Task Store_and_load_messages()
    {
        var store = new MemStore();

        var seq1 = await store.AppendAsync("foo", "Hello World"u8.ToArray(), default);
        var seq2 = await store.AppendAsync("bar", "Second"u8.ToArray(), default);

        seq1.ShouldBe((ulong)1);
        seq2.ShouldBe((ulong)2);

        var state = await store.GetStateAsync(default);
        state.Messages.ShouldBe((ulong)2);
        state.FirstSeq.ShouldBe((ulong)1);
        state.LastSeq.ShouldBe((ulong)2);

        var msg1 = await store.LoadAsync(1, default);
        msg1.ShouldNotBeNull();
        msg1!.Subject.ShouldBe("foo");
        msg1.Payload.ToArray().ShouldBe("Hello World"u8.ToArray());

        var msg2 = await store.LoadAsync(2, default);
        msg2.ShouldNotBeNull();
        msg2!.Subject.ShouldBe("bar");
    }

    // Go: TestMemStoreBasics server/memstore_test.go
    [Fact]
    public async Task Load_non_existent_returns_null()
    {
        var store = new MemStore();

        await store.AppendAsync("foo", "data"u8.ToArray(), default);

        (await store.LoadAsync(99, default)).ShouldBeNull();
        (await store.LoadAsync(0, default)).ShouldBeNull();
    }

    // Go: TestMemStoreEraseMsg server/memstore_test.go
    [Fact]
    public async Task Remove_messages()
    {
        var store = new MemStore();

        for (var i = 0; i < 5; i++)
            await store.AppendAsync("foo", Encoding.UTF8.GetBytes($"msg-{i}"), default);

        (await store.RemoveAsync(2, default)).ShouldBeTrue();
        (await store.RemoveAsync(4, default)).ShouldBeTrue();

        var state = await store.GetStateAsync(default);
        state.Messages.ShouldBe((ulong)3);

        (await store.LoadAsync(2, default)).ShouldBeNull();
        (await store.LoadAsync(4, default)).ShouldBeNull();
        (await store.LoadAsync(1, default)).ShouldNotBeNull();
        (await store.LoadAsync(3, default)).ShouldNotBeNull();
        (await store.LoadAsync(5, default)).ShouldNotBeNull();
    }

    // Go: TestMemStoreEraseMsg server/memstore_test.go
    [Fact]
    public async Task Remove_non_existent_returns_false()
    {
        var store = new MemStore();

        await store.AppendAsync("foo", "data"u8.ToArray(), default);

        (await store.RemoveAsync(99, default)).ShouldBeFalse();
    }

    // Go: TestMemStorePurge server/memstore_test.go
    [Fact]
    public async Task Purge_clears_all()
    {
        var store = new MemStore();

        for (var i = 0; i < 10; i++)
            await store.AppendAsync("foo", "data"u8.ToArray(), default);

        (await store.GetStateAsync(default)).Messages.ShouldBe((ulong)10);

        await store.PurgeAsync(default);

        var state = await store.GetStateAsync(default);
        state.Messages.ShouldBe((ulong)0);
        state.Bytes.ShouldBe((ulong)0);
    }

    // Go: TestMemStorePurge server/memstore_test.go
    [Fact]
    public async Task Purge_empty_store_is_safe()
    {
        var store = new MemStore();

        await store.PurgeAsync(default);

        (await store.GetStateAsync(default)).Messages.ShouldBe((ulong)0);
    }

    // Go: TestMemStoreTimeStamps server/memstore_test.go
    [Fact]
    public async Task Timestamps_non_decreasing()
    {
        var store = new MemStore();

        for (var i = 0; i < 10; i++)
            await store.AppendAsync("foo", "data"u8.ToArray(), default);

        var messages = await store.ListAsync(default);
        messages.Count.ShouldBe(10);

        DateTime? prev = null;
        foreach (var msg in messages)
        {
            if (prev.HasValue)
                msg.TimestampUtc.ShouldBeGreaterThanOrEqualTo(prev.Value);
            prev = msg.TimestampUtc;
        }
    }

    // Go: TestMemStoreMsgHeaders (adapted) server/memstore_test.go
    [Fact]
    public async Task Payload_with_header_bytes_round_trips()
    {
        var store = new MemStore();

        var headerBytes = "NATS/1.0\r\nName: derek\r\n\r\n"u8.ToArray();
        var bodyBytes = "Hello World"u8.ToArray();
        byte[] combined = [.. headerBytes, .. bodyBytes];

        await store.AppendAsync("foo", combined, default);

        var msg = await store.LoadAsync(1, default);
        msg.ShouldNotBeNull();
        msg!.Payload.ToArray().ShouldBe(combined);
    }

    // Go: TestMemStoreBasics server/memstore_test.go
    [Fact]
    public async Task LoadLastBySubject_returns_most_recent()
    {
        var store = new MemStore();

        await store.AppendAsync("foo", "first"u8.ToArray(), default);
        await store.AppendAsync("bar", "other"u8.ToArray(), default);
        await store.AppendAsync("foo", "second"u8.ToArray(), default);
        await store.AppendAsync("foo", "third"u8.ToArray(), default);

        var last = await store.LoadLastBySubjectAsync("foo", default);
        last.ShouldNotBeNull();
        last!.Payload.ToArray().ShouldBe("third"u8.ToArray());
        last.Sequence.ShouldBe((ulong)4);

        (await store.LoadLastBySubjectAsync("does.not.exist", default)).ShouldBeNull();
    }

    // Go: TestMemStoreMsgLimit server/memstore_test.go
    [Fact]
    public async Task TrimToMaxMessages_evicts_oldest()
    {
        var store = new MemStore();

        for (var i = 0; i < 20; i++)
            await store.AppendAsync("foo", Encoding.UTF8.GetBytes($"msg-{i}"), default);

        store.TrimToMaxMessages(10);

        var state = await store.GetStateAsync(default);
        state.Messages.ShouldBe((ulong)10);
        state.FirstSeq.ShouldBe((ulong)11);
        state.LastSeq.ShouldBe((ulong)20);

        (await store.LoadAsync(1, default)).ShouldBeNull();
        (await store.LoadAsync(10, default)).ShouldBeNull();
        (await store.LoadAsync(11, default)).ShouldNotBeNull();
    }

    // Go: TestMemStoreMsgLimit server/memstore_test.go
    [Fact]
    public async Task TrimToMaxMessages_to_zero()
    {
        var store = new MemStore();

        for (var i = 0; i < 5; i++)
            await store.AppendAsync("foo", "data"u8.ToArray(), default);

        store.TrimToMaxMessages(0);

        (await store.GetStateAsync(default)).Messages.ShouldBe((ulong)0);
    }

    // Go: TestMemStoreBytesLimit server/memstore_test.go
    [Fact]
    public async Task Bytes_tracks_payload_sizes()
    {
        var store = new MemStore();

        var payload = new byte[100];
        for (var i = 0; i < 5; i++)
            await store.AppendAsync("foo", payload, default);

        var state = await store.GetStateAsync(default);
        // Go parity: MsgSize = subj.Length + hdr + data + 16 overhead
        // "foo"(3) + 100 + 16 = 119 per msg × 5 = 595
        state.Bytes.ShouldBe((ulong)595);
    }

    // Go: TestMemStoreBytesLimit server/memstore_test.go
    [Fact]
    public async Task Bytes_decrease_after_remove()
    {
        var store = new MemStore();

        var payload = new byte[100];
        for (var i = 0; i < 5; i++)
            await store.AppendAsync("foo", payload, default);

        await store.RemoveAsync(1, default);
        await store.RemoveAsync(3, default);

        var state = await store.GetStateAsync(default);
        // Go parity: MsgSize = subj.Length + hdr + data + 16 overhead
        // "foo"(3) + 100 + 16 = 119 per msg × 3 remaining = 357
        state.Bytes.ShouldBe((ulong)357);
    }

    // Snapshot and restore.
    [Fact]
    public async Task Snapshot_and_restore()
    {
        var store = new MemStore();

        for (var i = 0; i < 20; i++)
            await store.AppendAsync("foo", Encoding.UTF8.GetBytes($"msg-{i}"), default);

        var snap = await store.CreateSnapshotAsync(default);
        snap.Length.ShouldBeGreaterThan(0);

        var restored = new MemStore();
        await restored.RestoreSnapshotAsync(snap, default);

        var srcState = await store.GetStateAsync(default);
        var dstState = await restored.GetStateAsync(default);
        dstState.Messages.ShouldBe(srcState.Messages);
        dstState.FirstSeq.ShouldBe(srcState.FirstSeq);
        dstState.LastSeq.ShouldBe(srcState.LastSeq);

        for (ulong i = 1; i <= srcState.Messages; i++)
        {
            var original = await store.LoadAsync(i, default);
            var copy = await restored.LoadAsync(i, default);
            copy.ShouldNotBeNull();
            copy!.Payload.ToArray().ShouldBe(original!.Payload.ToArray());
        }
    }

    // Snapshot after removes.
    [Fact]
    public async Task Snapshot_after_removes()
    {
        var store = new MemStore();

        for (var i = 0; i < 10; i++)
            await store.AppendAsync("foo", Encoding.UTF8.GetBytes($"msg-{i}"), default);

        await store.RemoveAsync(2, default);
        await store.RemoveAsync(5, default);
        await store.RemoveAsync(8, default);

        var snap = await store.CreateSnapshotAsync(default);

        var restored = new MemStore();
        await restored.RestoreSnapshotAsync(snap, default);

        var dstState = await restored.GetStateAsync(default);
        dstState.Messages.ShouldBe((ulong)7);

        (await restored.LoadAsync(2, default)).ShouldBeNull();
        (await restored.LoadAsync(5, default)).ShouldBeNull();
        (await restored.LoadAsync(8, default)).ShouldBeNull();
        (await restored.LoadAsync(1, default)).ShouldNotBeNull();
    }

    // ListAsync ordered.
    [Fact]
    public async Task ListAsync_returns_ordered()
    {
        var store = new MemStore();

        await store.AppendAsync("c", "three"u8.ToArray(), default);
        await store.AppendAsync("a", "one"u8.ToArray(), default);
        await store.AppendAsync("b", "two"u8.ToArray(), default);

        var messages = await store.ListAsync(default);
        messages.Count.ShouldBe(3);
        messages[0].Sequence.ShouldBe((ulong)1);
        messages[1].Sequence.ShouldBe((ulong)2);
        messages[2].Sequence.ShouldBe((ulong)3);
    }

    // Purge then append.
    [Fact]
    public async Task Purge_then_append()
    {
        var store = new MemStore();

        for (var i = 0; i < 5; i++)
            await store.AppendAsync("foo", "data"u8.ToArray(), default);

        await store.PurgeAsync(default);

        var seq = await store.AppendAsync("foo", "after purge"u8.ToArray(), default);
        seq.ShouldBeGreaterThan((ulong)0);

        var msg = await store.LoadAsync(seq, default);
        msg.ShouldNotBeNull();
        msg!.Payload.ToArray().ShouldBe("after purge"u8.ToArray());
    }

    // Empty payload.
    [Fact]
    public async Task Empty_payload_round_trips()
    {
        var store = new MemStore();

        await store.AppendAsync("foo", ReadOnlyMemory<byte>.Empty, default);

        var msg = await store.LoadAsync(1, default);
        msg.ShouldNotBeNull();
        msg!.Payload.Length.ShouldBe(0);
    }

    // State on empty store.
    [Fact]
    public async Task Empty_store_state()
    {
        var store = new MemStore();

        var state = await store.GetStateAsync(default);
        state.Messages.ShouldBe((ulong)0);
        state.Bytes.ShouldBe((ulong)0);
        state.FirstSeq.ShouldBe((ulong)0);
        state.LastSeq.ShouldBe((ulong)0);
    }
}