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

// Ported from golang/nats-server/server/memstore_test.go:
//   TestMemStoreBasics, TestMemStorePurge, TestMemStoreMsgHeaders (adapted),
//   TestMemStoreTimeStamps, TestMemStoreEraseMsg

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

namespace NATS.Server.Tests.JetStream.Storage;

public class MemStoreBasicTests
{
    // Go ref: TestMemStoreBasics — store a message, verify sequence, state, and payload round-trip.
    [Fact]
    public async Task Store_and_load_messages()
    {
        var store = new MemStore();

        var payload1 = "Hello World"u8.ToArray();
        var payload2 = "Second message"u8.ToArray();
        var payload3 = "Third message"u8.ToArray();
        var payload4 = "Fourth message"u8.ToArray();
        var payload5 = "Fifth message"u8.ToArray();

        var seq1 = await store.AppendAsync("foo", payload1, default);
        var seq2 = await store.AppendAsync("foo", payload2, default);
        var seq3 = await store.AppendAsync("bar", payload3, default);
        var seq4 = await store.AppendAsync("bar", payload4, default);
        var seq5 = await store.AppendAsync("baz", payload5, default);

        seq1.ShouldBe((ulong)1);
        seq2.ShouldBe((ulong)2);
        seq3.ShouldBe((ulong)3);
        seq4.ShouldBe((ulong)4);
        seq5.ShouldBe((ulong)5);

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

        var loaded1 = await store.LoadAsync(1, default);
        loaded1.ShouldNotBeNull();
        loaded1.Subject.ShouldBe("foo");
        loaded1.Sequence.ShouldBe((ulong)1);
        loaded1.Payload.Span.SequenceEqual(payload1).ShouldBeTrue();

        var loaded3 = await store.LoadAsync(3, default);
        loaded3.ShouldNotBeNull();
        loaded3.Subject.ShouldBe("bar");
        loaded3.Payload.Span.SequenceEqual(payload3).ShouldBeTrue();

        var loaded5 = await store.LoadAsync(5, default);
        loaded5.ShouldNotBeNull();
        loaded5.Subject.ShouldBe("baz");
        loaded5.Payload.Span.SequenceEqual(payload5).ShouldBeTrue();
    }

    // Go ref: TestMemStoreMsgHeaders (adapted) — MemStore stores and retrieves arbitrary payloads;
    // the .NET StoredMessage does not have a separate headers field (headers are embedded in the
    // payload by the protocol layer), so this test verifies that binary payload content round-trips
    // exactly including non-ASCII byte sequences that mimic header framing.
    [Fact]
    public async Task Store_preserves_payload_bytes_including_header_framing()
    {
        var store = new MemStore();

        // Simulate a payload that includes NATS header framing bytes followed by body bytes,
        // as the protocol layer would hand them to the store.
        var headerBytes = Encoding.ASCII.GetBytes("NATS/1.0\r\nName: derek\r\n\r\n");
        var bodyBytes = "Hello World"u8.ToArray();
        byte[] combined = [.. headerBytes, .. bodyBytes];

        var seq = await store.AppendAsync("foo", combined, default);
        seq.ShouldBe((ulong)1);

        var loaded = await store.LoadAsync(1, default);
        loaded.ShouldNotBeNull();
        loaded.Subject.ShouldBe("foo");
        loaded.Payload.Length.ShouldBe(combined.Length);
        loaded.Payload.Span.SequenceEqual(combined).ShouldBeTrue();
    }

    // Go ref: TestMemStoreEraseMsg — remove a message returns true; subsequent load returns null.
    [Fact]
    public async Task Remove_messages_updates_state()
    {
        var store = new MemStore();

        var seq1 = await store.AppendAsync("foo", "one"u8.ToArray(), default);
        var seq2 = await store.AppendAsync("foo", "two"u8.ToArray(), default);
        var seq3 = await store.AppendAsync("foo", "three"u8.ToArray(), default);
        var seq4 = await store.AppendAsync("foo", "four"u8.ToArray(), default);
        var seq5 = await store.AppendAsync("foo", "five"u8.ToArray(), default);

        var stateBefore = await store.GetStateAsync(default);
        stateBefore.Messages.ShouldBe((ulong)5);

        // Remove seq2 and seq4 (interior messages).
        (await store.RemoveAsync(seq2, default)).ShouldBeTrue();
        (await store.RemoveAsync(seq4, default)).ShouldBeTrue();

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

        // Removed sequences are no longer loadable.
        (await store.LoadAsync(seq2, default)).ShouldBeNull();
        (await store.LoadAsync(seq4, default)).ShouldBeNull();

        // Remaining messages are still loadable.
        (await store.LoadAsync(seq1, default)).ShouldNotBeNull();
        (await store.LoadAsync(seq3, default)).ShouldNotBeNull();
        (await store.LoadAsync(seq5, default)).ShouldNotBeNull();

        // Removing a non-existent sequence returns false.
        (await store.RemoveAsync(99, default)).ShouldBeFalse();
    }

    // Go ref: TestMemStorePurge — purge clears all messages and resets state.
    [Fact]
    public async Task Purge_clears_all_messages()
    {
        var store = new MemStore();

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

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

        await store.PurgeAsync(default);

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

    // Go ref: TestMemStoreTimeStamps — each stored message gets a distinct, monotonically
    // increasing timestamp.
    [Fact]
    public async Task Stored_messages_have_distinct_non_decreasing_timestamps()
    {
        var store = new MemStore();
        const int count = 5;

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

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

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

    // Go ref: TestMemStoreBasics — LoadLastBySubject returns the highest-sequence message
    // for the given subject.
    [Fact]
    public async Task Load_last_by_subject_returns_most_recent_for_that_subject()
    {
        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.Span.SequenceEqual("third"u8).ShouldBeTrue();
        last.Subject.ShouldBe("foo");

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