natsdotnet/tests/NATS.Server.Raft.Tests/Raft/RaftPreVoteTests.cs

using NATS.Server.Raft;

namespace NATS.Server.Raft.Tests.Raft;

/// <summary>
/// Tests for B6: Pre-Vote Protocol.
/// Go reference: raft.go:1600-1700 (pre-vote logic).
/// Pre-vote prevents partitioned nodes from disrupting the cluster by
/// incrementing their term without actually winning an election.
/// </summary>
public class RaftPreVoteTests
{
    // -- Helpers --

    private static (RaftNode[] nodes, InMemoryRaftTransport transport) CreateCluster(int size)
    {
        var transport = new InMemoryRaftTransport();
        var nodes = Enumerable.Range(1, size)
            .Select(i => new RaftNode($"n{i}", transport))
            .ToArray();
        foreach (var node in nodes)
        {
            transport.Register(node);
            node.ConfigureCluster(nodes);
        }

        return (nodes, transport);
    }

    private static RaftNode ElectLeader(RaftNode[] nodes)
    {
        var candidate = nodes[0];
        candidate.StartElection(nodes.Length);
        foreach (var voter in nodes.Skip(1))
            candidate.ReceiveVote(voter.GrantVote(candidate.Term, candidate.Id), nodes.Length);
        return candidate;
    }

    // -- Wire format tests --

    [Fact]
    public void PreVote_request_encoding_roundtrip()
    {
        var request = new RaftPreVoteRequestWire(
            Term: 5,
            LastTerm: 4,
            LastIndex: 100,
            CandidateId: "n1");

        var encoded = request.Encode();
        encoded.Length.ShouldBe(RaftWireConstants.VoteRequestLen); // 32 bytes

        var decoded = RaftPreVoteRequestWire.Decode(encoded);
        decoded.Term.ShouldBe(5UL);
        decoded.LastTerm.ShouldBe(4UL);
        decoded.LastIndex.ShouldBe(100UL);
        decoded.CandidateId.ShouldBe("n1");
    }

    [Fact]
    public void PreVote_response_encoding_roundtrip()
    {
        var response = new RaftPreVoteResponseWire(
            Term: 5,
            PeerId: "n2",
            Granted: true);

        var encoded = response.Encode();
        encoded.Length.ShouldBe(RaftWireConstants.VoteResponseLen); // 17 bytes

        var decoded = RaftPreVoteResponseWire.Decode(encoded);
        decoded.Term.ShouldBe(5UL);
        decoded.PeerId.ShouldBe("n2");
        decoded.Granted.ShouldBeTrue();
    }

    [Fact]
    public void PreVote_response_denied_roundtrip()
    {
        var response = new RaftPreVoteResponseWire(Term: 3, PeerId: "n3", Granted: false);
        var decoded = RaftPreVoteResponseWire.Decode(response.Encode());
        decoded.Granted.ShouldBeFalse();
        decoded.PeerId.ShouldBe("n3");
        decoded.Term.ShouldBe(3UL);
    }

    [Fact]
    public void PreVote_request_decode_throws_on_wrong_length()
    {
        Should.Throw<ArgumentException>(() =>
            RaftPreVoteRequestWire.Decode(new byte[10]));
    }

    [Fact]
    public void PreVote_response_decode_throws_on_wrong_length()
    {
        Should.Throw<ArgumentException>(() =>
            RaftPreVoteResponseWire.Decode(new byte[10]));
    }

    // -- RequestPreVote logic tests --

    [Fact]
    public void PreVote_granted_when_candidate_log_is_up_to_date()
    {
        // Go reference: raft.go pre-vote grants when candidate log >= voter log
        var node = new RaftNode("voter");
        node.Log.Append(1, "cmd-1"); // voter has entry at index 1, term 1

        // Candidate has same term and same or higher index: should grant
        var granted = node.RequestPreVote(
            term: (ulong)node.Term,
            lastTerm: 1,
            lastIndex: 1,
            candidateId: "candidate");
        granted.ShouldBeTrue();
    }

    [Fact]
    public void PreVote_granted_when_candidate_has_higher_term_log()
    {
        var node = new RaftNode("voter");
        node.Log.Append(1, "cmd-1"); // voter: term 1, index 1

        // Candidate has higher last term: should grant
        var granted = node.RequestPreVote(
            term: 0,
            lastTerm: 2,
            lastIndex: 1,
            candidateId: "candidate");
        granted.ShouldBeTrue();
    }

    [Fact]
    public void PreVote_denied_when_candidate_log_is_stale()
    {
        // Go reference: raft.go pre-vote denied when candidate log behind voter
        var node = new RaftNode("voter");
        node.TermState.CurrentTerm = 2;
        node.Log.Append(2, "cmd-1");
        node.Log.Append(2, "cmd-2"); // voter: term 2, index 2

        // Candidate has lower last term: should deny
        var granted = node.RequestPreVote(
            term: 2,
            lastTerm: 1,
            lastIndex: 5,
            candidateId: "candidate");
        granted.ShouldBeFalse();
    }

    [Fact]
    public void PreVote_denied_when_candidate_term_behind()
    {
        var node = new RaftNode("voter");
        node.TermState.CurrentTerm = 5;

        // Candidate's term is behind the voter's current term
        var granted = node.RequestPreVote(
            term: 3,
            lastTerm: 3,
            lastIndex: 100,
            candidateId: "candidate");
        granted.ShouldBeFalse();
    }

    [Fact]
    public void PreVote_granted_for_empty_logs()
    {
        // Both node and candidate have empty logs: grant
        var node = new RaftNode("voter");

        var granted = node.RequestPreVote(
            term: 0,
            lastTerm: 0,
            lastIndex: 0,
            candidateId: "candidate");
        granted.ShouldBeTrue();
    }

    // -- Pre-vote integration with election flow --

    [Fact]
    public void Successful_prevote_leads_to_real_election()
    {
        // Go reference: after pre-vote success, proceed to real election with term increment
        var (nodes, _) = CreateCluster(3);
        var candidate = nodes[0];
        var termBefore = candidate.Term;

        // With pre-vote enabled, CampaignWithPreVote should succeed (all peers have equal logs)
        // and then start a real election (incrementing term)
        candidate.PreVoteEnabled = true;
        candidate.CampaignWithPreVote();

        // Term should have been incremented by the real election
        candidate.Term.ShouldBe(termBefore + 1);
        candidate.Role.ShouldBe(RaftRole.Candidate);
    }

    [Fact]
    public void Failed_prevote_does_not_increment_term()
    {
        // Go reference: failed pre-vote stays follower, doesn't disrupt cluster
        var (nodes, _) = CreateCluster(3);
        var candidate = nodes[0];

        // Give the other nodes higher-term logs so pre-vote will be denied
        nodes[1].TermState.CurrentTerm = 10;
        nodes[1].Log.Append(10, "advanced-cmd");
        nodes[2].TermState.CurrentTerm = 10;
        nodes[2].Log.Append(10, "advanced-cmd");

        var termBefore = candidate.Term;
        candidate.PreVoteEnabled = true;
        candidate.CampaignWithPreVote();

        // Term should NOT have been incremented — pre-vote failed
        candidate.Term.ShouldBe(termBefore);
        candidate.Role.ShouldBe(RaftRole.Follower);
    }

    [Fact]
    public void PreVote_disabled_goes_directly_to_election()
    {
        // When PreVoteEnabled is false, skip pre-vote and go straight to election
        var (nodes, _) = CreateCluster(3);
        var candidate = nodes[0];
        var termBefore = candidate.Term;

        candidate.PreVoteEnabled = false;
        candidate.CampaignWithPreVote();

        // Should have gone directly to election, incrementing term
        candidate.Term.ShouldBe(termBefore + 1);
        candidate.Role.ShouldBe(RaftRole.Candidate);
    }

    [Fact]
    public void Partitioned_node_with_stale_term_does_not_disrupt_via_prevote()
    {
        // Go reference: pre-vote prevents partitioned nodes from disrupting the cluster.
        // A node with a stale term that reconnects should fail the pre-vote round
        // and NOT increment its term, which would force other nodes to step down.
        var (nodes, _) = CreateCluster(3);

        // Simulate: n1 was partitioned and has term 0, others advanced to term 5
        nodes[1].TermState.CurrentTerm = 5;
        nodes[1].Log.Append(5, "cmd-a");
        nodes[1].Log.Append(5, "cmd-b");
        nodes[2].TermState.CurrentTerm = 5;
        nodes[2].Log.Append(5, "cmd-a");
        nodes[2].Log.Append(5, "cmd-b");

        var partitioned = nodes[0];
        partitioned.PreVoteEnabled = true;
        var termBefore = partitioned.Term;

        // Pre-vote should fail because the partitioned node has a stale log
        partitioned.CampaignWithPreVote();

        // The partitioned node should NOT have incremented its term
        partitioned.Term.ShouldBe(termBefore);
        partitioned.Role.ShouldBe(RaftRole.Follower);
    }

    [Fact]
    public void PreVote_enabled_by_default()
    {
        var node = new RaftNode("n1");
        node.PreVoteEnabled.ShouldBeTrue();
    }

    [Fact]
    public void StartPreVote_returns_true_when_majority_grants()
    {
        // All nodes have empty, equal logs: pre-vote should succeed
        var (nodes, _) = CreateCluster(3);
        var candidate = nodes[0];

        var result = candidate.StartPreVote();
        result.ShouldBeTrue();
    }

    [Fact]
    public void StartPreVote_returns_false_when_majority_denies()
    {
        var (nodes, _) = CreateCluster(3);
        var candidate = nodes[0];

        // Make majority have more advanced logs
        nodes[1].TermState.CurrentTerm = 10;
        nodes[1].Log.Append(10, "cmd");
        nodes[2].TermState.CurrentTerm = 10;
        nodes[2].Log.Append(10, "cmd");

        var result = candidate.StartPreVote();
        result.ShouldBeFalse();
    }
}