natsdotnet/tests/NATS.Server.JetStream.Tests/JetStream/Cluster/StreamRaftGroupTests.cs

// Go parity: golang/nats-server/server/jetstream_cluster.go
// Covers: Per-stream RAFT group message proposals, message count tracking,
// sequence tracking, leader change events, replica status reporting,
// and non-leader rejection.
using NATS.Server.JetStream.Cluster;

namespace NATS.Server.JetStream.Tests.JetStream.Cluster;

/// <summary>
/// Tests for StreamReplicaGroup stream-specific RAFT apply logic:
/// message proposals, message count, last sequence, leader change
/// event, and replica status reporting.
/// Go reference: jetstream_cluster.go processStreamMsg, processStreamEntries.
/// </summary>
public class StreamRaftGroupTests
{
    // ---------------------------------------------------------------
    // ProposeMessageAsync succeeds as leader
    // Go reference: jetstream_cluster.go processStreamMsg
    // ---------------------------------------------------------------

    [Fact]
    public async Task Propose_message_succeeds_as_leader()
    {
        var group = new StreamReplicaGroup("MSGS", replicas: 3);

        var index = await group.ProposeMessageAsync(
            "orders.new", ReadOnlyMemory<byte>.Empty, "hello"u8.ToArray(), default);

        index.ShouldBeGreaterThan(0);
    }

    // ---------------------------------------------------------------
    // ProposeMessageAsync fails when not leader
    // Go reference: jetstream_cluster.go leader check
    // ---------------------------------------------------------------

    [Fact]
    public async Task Propose_message_fails_when_not_leader()
    {
        var group = new StreamReplicaGroup("NOLEAD", replicas: 3);

        // Step down so the current leader is no longer leader
        group.Leader.RequestStepDown();

        await Should.ThrowAsync<InvalidOperationException>(async () =>
            await group.ProposeMessageAsync(
                "test.sub", ReadOnlyMemory<byte>.Empty, "data"u8.ToArray(), default));
    }

    // ---------------------------------------------------------------
    // Message count increments after proposal
    // Go reference: stream.go state.Msgs tracking
    // ---------------------------------------------------------------

    [Fact]
    public async Task Message_count_increments_after_proposal()
    {
        var group = new StreamReplicaGroup("COUNT", replicas: 3);

        group.MessageCount.ShouldBe(0);

        await group.ProposeMessageAsync("a.1", ReadOnlyMemory<byte>.Empty, "m1"u8.ToArray(), default);
        group.MessageCount.ShouldBe(1);

        await group.ProposeMessageAsync("a.2", ReadOnlyMemory<byte>.Empty, "m2"u8.ToArray(), default);
        group.MessageCount.ShouldBe(2);

        await group.ProposeMessageAsync("a.3", ReadOnlyMemory<byte>.Empty, "m3"u8.ToArray(), default);
        group.MessageCount.ShouldBe(3);
    }

    // ---------------------------------------------------------------
    // Last sequence tracks correctly
    // Go reference: stream.go state.LastSeq
    // ---------------------------------------------------------------

    [Fact]
    public async Task Last_sequence_tracks_correctly()
    {
        var group = new StreamReplicaGroup("SEQ", replicas: 3);

        group.LastSequence.ShouldBe(0);

        var idx1 = await group.ProposeMessageAsync("s.1", ReadOnlyMemory<byte>.Empty, "d1"u8.ToArray(), default);
        group.LastSequence.ShouldBe(idx1);

        var idx2 = await group.ProposeMessageAsync("s.2", ReadOnlyMemory<byte>.Empty, "d2"u8.ToArray(), default);
        group.LastSequence.ShouldBe(idx2);

        idx2.ShouldBeGreaterThan(idx1);
    }

    // ---------------------------------------------------------------
    // Step down triggers leader change event
    // Go reference: jetstream_cluster.go leader change notification
    // ---------------------------------------------------------------

    [Fact]
    public async Task Step_down_triggers_leader_change_event()
    {
        var group = new StreamReplicaGroup("EVENT", replicas: 3);
        var previousId = group.Leader.Id;

        LeaderChangedEventArgs? receivedArgs = null;
        group.LeaderChanged += (_, args) => receivedArgs = args;

        await group.StepDownAsync(default);

        receivedArgs.ShouldNotBeNull();
        receivedArgs.PreviousLeaderId.ShouldBe(previousId);
        receivedArgs.NewLeaderId.ShouldNotBe(previousId);
        receivedArgs.NewTerm.ShouldBeGreaterThan(0);
    }

    [Fact]
    public async Task Multiple_stepdowns_fire_leader_changed_each_time()
    {
        var group = new StreamReplicaGroup("MULTI_EVENT", replicas: 3);
        var eventCount = 0;
        group.LeaderChanged += (_, _) => eventCount++;

        await group.StepDownAsync(default);
        await group.StepDownAsync(default);
        await group.StepDownAsync(default);

        eventCount.ShouldBe(3);
    }

    // ---------------------------------------------------------------
    // Replica status reports correct state
    // Go reference: jetstream_cluster.go stream replica status
    // ---------------------------------------------------------------

    [Fact]
    public async Task Replica_status_reports_correct_state()
    {
        var group = new StreamReplicaGroup("STATUS", replicas: 3);

        await group.ProposeMessageAsync("x.1", ReadOnlyMemory<byte>.Empty, "m1"u8.ToArray(), default);
        await group.ProposeMessageAsync("x.2", ReadOnlyMemory<byte>.Empty, "m2"u8.ToArray(), default);

        var status = group.GetStatus();

        status.StreamName.ShouldBe("STATUS");
        status.LeaderId.ShouldBe(group.Leader.Id);
        status.LeaderTerm.ShouldBeGreaterThan(0);
        status.MessageCount.ShouldBe(2);
        status.LastSequence.ShouldBeGreaterThan(0);
        status.ReplicaCount.ShouldBe(3);
    }

    [Fact]
    public void Initial_status_has_zero_messages()
    {
        var group = new StreamReplicaGroup("EMPTY", replicas: 1);

        var status = group.GetStatus();

        status.MessageCount.ShouldBe(0);
        status.LastSequence.ShouldBe(0);
        status.ReplicaCount.ShouldBe(1);
    }

    // ---------------------------------------------------------------
    // Status updates after step down
    // ---------------------------------------------------------------

    [Fact]
    public async Task Status_reflects_new_leader_after_stepdown()
    {
        var group = new StreamReplicaGroup("NEWLEAD", replicas: 3);
        var statusBefore = group.GetStatus();

        await group.StepDownAsync(default);

        var statusAfter = group.GetStatus();
        statusAfter.LeaderId.ShouldNotBe(statusBefore.LeaderId);
    }

    // ---------------------------------------------------------------
    // ProposeAsync still works after ProposeMessageAsync
    // ---------------------------------------------------------------

    [Fact]
    public async Task ProposeAsync_and_ProposeMessageAsync_coexist()
    {
        var group = new StreamReplicaGroup("COEXIST", replicas: 3);

        var idx1 = await group.ProposeAsync("PUB test.1", default);
        var idx2 = await group.ProposeMessageAsync("test.2", ReadOnlyMemory<byte>.Empty, "data"u8.ToArray(), default);

        idx2.ShouldBeGreaterThan(idx1);
        group.MessageCount.ShouldBe(1); // Only ProposeMessageAsync increments message count
    }
}