natsnet/dotnet/tests/ZB.MOM.NatsNet.Server.Tests/JetStream/JetStreamBatchingCoreTests.cs

using NSubstitute;
using Shouldly;

namespace ZB.MOM.NatsNet.Server.Tests.JetStream;

public sealed class JetStreamBatchingCoreTests
{
    [Fact]
    public void GetBatchStoreDir_ValidInputs_ReturnsHashedBatchPath()
    {
        var storeDir = Path.Combine(Path.GetTempPath(), $"jsa-{Guid.NewGuid():N}");
        var (batchName, batchPath) = JetStreamBatching.GetBatchStoreDir(storeDir, "ORDERS", "batch-A");

        batchName.ShouldBe(NatsServer.GetHash("batch-A"));
        batchPath.ShouldBe(Path.Combine(storeDir, "_streams", "ORDERS", "batches", batchName));
    }

    [Fact]
    public void NewBatchStore_FileSingleReplica_CreatesFileStore()
    {
        var tempRoot = Path.Combine(Path.GetTempPath(), $"batch-store-{Guid.NewGuid():N}");
        var stream = new StreamConfig
        {
            Name = "ORDERS",
            Replicas = 1,
            Storage = StorageType.FileStorage,
        };

        var store = JetStreamBatching.NewBatchStore(tempRoot, stream, "batch-A");

        try
        {
            store.ShouldBeOfType<JetStreamFileStore>();
        }
        finally
        {
            store.Stop();
        }
    }

    [Fact]
    public void NewBatchStore_MemoryOrReplicated_CreatesMemStore()
    {
        var tempRoot = Path.Combine(Path.GetTempPath(), $"batch-store-{Guid.NewGuid():N}");
        var stream = new StreamConfig
        {
            Name = "ORDERS",
            Replicas = 3,
            Storage = StorageType.FileStorage,
        };

        var store = JetStreamBatching.NewBatchStore(tempRoot, stream, "batch-A");

        try
        {
            store.ShouldBeOfType<JetStreamMemStore>();
        }
        finally
        {
            store.Stop();
        }
    }

    [Fact]
    public void ReadyForCommit_TimerStopped_FlushesAndReturnsTrue()
    {
        var store = Substitute.For<IStreamStore>();
        var timer = Substitute.For<IBatchTimer>();
        timer.Stop().Returns(true);

        var group = new BatchGroup
        {
            Store = store,
            BatchTimer = timer,
        };

        group.ReadyForCommit().ShouldBeTrue();
        store.Received(1).FlushAllPending();
    }

    [Fact]
    public void ReadyForCommit_TimerAlreadyExpired_DoesNotFlushAndReturnsFalse()
    {
        var store = Substitute.For<IStreamStore>();
        var timer = Substitute.For<IBatchTimer>();
        timer.Stop().Returns(false);

        var group = new BatchGroup
        {
            Store = store,
            BatchTimer = timer,
        };

        group.ReadyForCommit().ShouldBeFalse();
        store.DidNotReceive().FlushAllPending();
    }

    [Fact]
    public void CleanupLocked_BatchPresent_StopsTimerDeletesStoreRemovesGroupAndDecrementsGlobalInflight()
    {
        JetStreamBatching.ResetGlobalInflightBatchesForTest();
        JetStreamBatching.IncrementGlobalInflightBatchesForTest();

        var store = Substitute.For<IStreamStore>();
        var timer = Substitute.For<IBatchTimer>();
        timer.Stop().Returns(true);

        var group = new BatchGroup
        {
            Store = store,
            BatchTimer = timer,
        };

        var batching = new Batching();
        batching.Group["batch-A"] = group;

        group.CleanupLocked("batch-A", batching);

        timer.Received(1).Stop();
        store.Received(1).Delete(true);
        batching.Group.ContainsKey("batch-A").ShouldBeFalse();
        JetStreamBatching.GlobalInflightBatches.ShouldBe(0);
    }

    [Fact]
    public void StopLocked_BatchPresent_StopsTimerStopsStoreAndDecrementsGlobalInflight()
    {
        JetStreamBatching.ResetGlobalInflightBatchesForTest();
        JetStreamBatching.IncrementGlobalInflightBatchesForTest();

        var store = Substitute.For<IStreamStore>();
        var timer = Substitute.For<IBatchTimer>();

        var group = new BatchGroup
        {
            Store = store,
            BatchTimer = timer,
        };

        group.StopLocked();

        timer.Received(1).Stop();
        store.Received(1).Stop();
        JetStreamBatching.GlobalInflightBatches.ShouldBe(0);
    }
}