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

using NATS.Server.Raft;

namespace NATS.Server.Raft.Tests;

public class RaftTransportPersistenceTests
{
    [Fact]
    public async Task Raft_node_recovers_log_and_term_after_restart()
    {
        await using var fx = await RaftFixture.StartPersistentClusterAsync();
        var idx = await fx.Leader.ProposeAsync("cmd", default);
        await fx.RestartNodeAsync("n2");
        (await fx.ReadNodeAppliedIndexAsync("n2")).ShouldBeGreaterThanOrEqualTo(idx);
    }
}

internal sealed class RaftFixture : IAsyncDisposable
{
    private readonly string _root;
    private readonly InMemoryRaftTransport _transport;
    private readonly Dictionary<string, RaftNode> _nodes;

    private RaftFixture(string root, InMemoryRaftTransport transport, Dictionary<string, RaftNode> nodes)
    {
        _root = root;
        _transport = transport;
        _nodes = nodes;
    }

    public RaftNode Leader => _nodes["n1"];

    public static Task<RaftFixture> StartPersistentClusterAsync()
    {
        var root = Path.Combine(Path.GetTempPath(), $"nats-raft-{Guid.NewGuid():N}");
        Directory.CreateDirectory(root);

        var transport = new InMemoryRaftTransport();
        var nodes = new Dictionary<string, RaftNode>(StringComparer.Ordinal);
        foreach (var id in new[] { "n1", "n2", "n3" })
        {
            var node = new RaftNode(id, transport, Path.Combine(root, id));
            transport.Register(node);
            nodes[id] = node;
        }

        var all = nodes.Values.ToArray();
        foreach (var node in all)
            node.ConfigureCluster(all);

        var leader = nodes["n1"];
        leader.StartElection(all.Length);
        leader.ReceiveVote(nodes["n2"].GrantVote(leader.Term), all.Length);
        leader.ReceiveVote(nodes["n3"].GrantVote(leader.Term), all.Length);

        return Task.FromResult(new RaftFixture(root, transport, nodes));
    }

    public async Task RestartNodeAsync(string id)
    {
        var nodePath = Path.Combine(_root, id);
        var replacement = new RaftNode(id, _transport, nodePath);
        await replacement.LoadPersistedStateAsync(default);
        _transport.Register(replacement);
        _nodes[id] = replacement;

        var all = _nodes.Values.ToArray();
        foreach (var node in all)
            node.ConfigureCluster(all);
    }

    public Task<long> ReadNodeAppliedIndexAsync(string id)
    {
        return Task.FromResult(_nodes[id].AppliedIndex);
    }

    public ValueTask DisposeAsync()
    {
        try
        {
            if (Directory.Exists(_root))
                Directory.Delete(_root, recursive: true);
        }
        catch (IOException ex)
        {
            System.Diagnostics.Debug.WriteLine($"Failed to clean up temp directory {_root}: {ex.Message}");
        }

        return ValueTask.CompletedTask;
    }
}