natsdotnet/tests/NATS.E2E.Cluster.Tests/Infrastructure/HubLeafFixture.cs

using System.Net.Http.Json;
using System.Text.Json.Serialization;
using NATS.Client.Core;

namespace NATS.E2E.Cluster.Tests.Infrastructure;

public sealed class HubLeafFixture : IAsyncLifetime
{
    private NatsServerProcess? _hub;
    private NatsServerProcess? _leaf;

    // Pre-allocated ports stored for restart reuse
    private readonly int _hubPort;
    private readonly int _leafPort;
    private readonly int _leafListenPort;
    private readonly int _hubMonitorPort;

    // Configs stored for restart reuse
    private readonly string _hubConfig;
    private readonly string _leafConfig;

    public int HubPort => _hubPort;
    public int LeafPort => _leafPort;

    public HubLeafFixture()
    {
        _hubPort = NatsServerProcess.AllocateFreePort();
        _leafPort = NatsServerProcess.AllocateFreePort();
        _leafListenPort = NatsServerProcess.AllocateFreePort();
        _hubMonitorPort = NatsServerProcess.AllocateFreePort();

        _hubConfig = $$"""
            server_name: hub
            leafnodes {
                listen: 127.0.0.1:{{_leafListenPort}}
            }
            """;

        _leafConfig = $$"""
            server_name: leaf
            leafnodes {
                remotes [
                    { url: "nats-leaf://127.0.0.1:{{_leafListenPort}}" }
                ]
            }
            """;
    }

    public async Task InitializeAsync()
    {
        _hub = new NatsServerProcess(_hubPort, configContent: _hubConfig, enableMonitoring: true, monitorPort: _hubMonitorPort);
        await _hub.StartAsync();

        _leaf = new NatsServerProcess(_leafPort, configContent: _leafConfig);
        await _leaf.StartAsync();

        await WaitForLeafConnectionAsync();
    }

    /// <summary>
    /// Kills the server at the given index (0 = hub, 1 = leaf).
    /// </summary>
    public async Task KillNode(int index)
    {
        if (index == 0 && _hub is not null)
        {
            await _hub.DisposeAsync();
            _hub = null;
        }
        else if (index == 1 && _leaf is not null)
        {
            await _leaf.DisposeAsync();
            _leaf = null;
        }
        else
        {
            throw new ArgumentOutOfRangeException(nameof(index), "Index must be 0 (hub) or 1 (leaf).");
        }
    }

    /// <summary>
    /// Restarts the server at the given index on its original ports.
    /// For hub (0): restarts hub and waits for TCP ready.
    /// For leaf (1): restarts leaf and waits for TCP ready.
    /// After either restart, call WaitForLeafConnectionAsync separately.
    /// </summary>
    public async Task RestartNode(int index)
    {
        if (index == 0)
        {
            if (_hub is not null)
                await _hub.DisposeAsync();

            _hub = new NatsServerProcess(_hubPort, configContent: _hubConfig, enableMonitoring: true, monitorPort: _hubMonitorPort);
            await _hub.StartAsync();
        }
        else if (index == 1)
        {
            if (_leaf is not null)
                await _leaf.DisposeAsync();

            _leaf = new NatsServerProcess(_leafPort, configContent: _leafConfig);
            await _leaf.StartAsync();
        }
        else
        {
            throw new ArgumentOutOfRangeException(nameof(index), "Index must be 0 (hub) or 1 (leaf).");
        }
    }

    /// <summary>
    /// Polls hub's /leafz until NumLeafs >= 1. 30s timeout, 200ms polling.
    /// Handles hub being down gracefully by retrying on network errors.
    /// </summary>
    public async Task WaitForLeafConnectionAsync()
    {
        using var http = new HttpClient { Timeout = TimeSpan.FromSeconds(2) };
        using var deadline = new CancellationTokenSource(TimeSpan.FromSeconds(30));
        using var timer = new PeriodicTimer(TimeSpan.FromMilliseconds(200));

        while (await timer.WaitForNextTickAsync(deadline.Token).ConfigureAwait(false))
        {
            try
            {
                var leafz = await http.GetFromJsonAsync<Leafz>(
                    $"http://127.0.0.1:{_hubMonitorPort}/leafz",
                    deadline.Token);

                if (leafz?.NumLeafs >= 1)
                    return;
            }
            catch (Exception ex) when (ex is HttpRequestException or TaskCanceledException
                                        && !deadline.IsCancellationRequested)
            {
                // Hub not yet ready or per-request timeout — retry on next tick
                _ = ex;
            }
        }

        throw new TimeoutException("Leaf node did not connect to hub within 30s.");
    }

    public NatsConnection CreateHubClient()
        => new(new NatsOpts { Url = $"nats://127.0.0.1:{_hubPort}" });

    public NatsConnection CreateLeafClient()
        => new(new NatsOpts { Url = $"nats://127.0.0.1:{_leafPort}" });

    public async Task DisposeAsync()
    {
        if (_leaf is not null)
            await _leaf.DisposeAsync();

        if (_hub is not null)
            await _hub.DisposeAsync();
    }

    private sealed class Leafz
    {
        [JsonPropertyName("num_leafs")]
        public int NumLeafs { get; init; }
    }
}

[CollectionDefinition("E2E-HubLeaf")]
public class HubLeafCollection : ICollectionFixture<HubLeafFixture>;