natsdotnet/tests/NATS.Server.LeafNodes.Tests/LeafProtocolTests.cs

using System.Net;
using System.Net.Sockets;
using System.Text;
using Microsoft.Extensions.Logging.Abstractions;
using NATS.Server.Configuration;
using NATS.Server.TestUtilities;

namespace NATS.Server.LeafNodes.Tests;

public class LeafProtocolTests
{
    [Fact]
    public async Task Leaf_link_propagates_subscription_and_message_flow()
    {
        await using var fx = await LeafProtocolTestFixture.StartHubSpokeAsync();
        await fx.SubscribeSpokeAsync("leaf.>");
        await fx.PublishHubAsync("leaf.msg", "x");
        (await fx.ReadSpokeMessageAsync()).ShouldContain("x");
    }
}

internal sealed class LeafProtocolTestFixture : IAsyncDisposable
{
    private readonly NatsServer _hub;
    private readonly NatsServer _spoke;
    private readonly CancellationTokenSource _hubCts;
    private readonly CancellationTokenSource _spokeCts;
    private Socket? _spokeSubscriber;
    private Socket? _hubPublisher;

    private LeafProtocolTestFixture(NatsServer hub, NatsServer spoke, CancellationTokenSource hubCts, CancellationTokenSource spokeCts)
    {
        _hub = hub;
        _spoke = spoke;
        _hubCts = hubCts;
        _spokeCts = spokeCts;
    }

    public static async Task<LeafProtocolTestFixture> StartHubSpokeAsync()
    {
        var hubOptions = new NatsOptions
        {
            Host = "127.0.0.1",
            Port = 0,
            LeafNode = new LeafNodeOptions
            {
                Host = "127.0.0.1",
                Port = 0,
            },
        };

        var hub = new NatsServer(hubOptions, NullLoggerFactory.Instance);
        var hubCts = new CancellationTokenSource();
        _ = hub.StartAsync(hubCts.Token);
        await hub.WaitForReadyAsync();

        var spokeOptions = new NatsOptions
        {
            Host = "127.0.0.1",
            Port = 0,
            LeafNode = new LeafNodeOptions
            {
                Host = "127.0.0.1",
                Port = 0,
                Remotes = [hub.LeafListen!],
            },
        };

        var spoke = new NatsServer(spokeOptions, NullLoggerFactory.Instance);
        var spokeCts = new CancellationTokenSource();
        _ = spoke.StartAsync(spokeCts.Token);
        await spoke.WaitForReadyAsync();

        await PollHelper.WaitUntilAsync(() => hub.Stats.Leafs > 0 && spoke.Stats.Leafs > 0);

        return new LeafProtocolTestFixture(hub, spoke, hubCts, spokeCts);
    }

    public async Task SubscribeSpokeAsync(string subject)
    {
        var sock = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await sock.ConnectAsync(IPAddress.Loopback, _spoke.Port);
        _spokeSubscriber = sock;

        _ = await ReadLineAsync(sock); // INFO
        await sock.SendAsync(Encoding.ASCII.GetBytes($"CONNECT {{}}\r\nSUB {subject} 1\r\nPING\r\n"));
        await SocketTestHelper.ReadUntilAsync(sock, "PONG");
    }

    public async Task PublishHubAsync(string subject, string payload)
    {
        var sock = _hubPublisher;
        if (sock == null)
        {
            sock = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
            await sock.ConnectAsync(IPAddress.Loopback, _hub.Port);
            _hubPublisher = sock;
            _ = await ReadLineAsync(sock); // INFO
            await sock.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPING\r\n"));
            await SocketTestHelper.ReadUntilAsync(sock, "PONG");
        }

        await sock.SendAsync(Encoding.ASCII.GetBytes($"PUB {subject} {payload.Length}\r\n{payload}\r\nPING\r\n"));
        await SocketTestHelper.ReadUntilAsync(sock, "PONG");
    }

    public Task<string> ReadSpokeMessageAsync()
    {
        if (_spokeSubscriber == null)
            throw new InvalidOperationException("Spoke subscriber was not initialized.");

        return SocketTestHelper.ReadUntilAsync(_spokeSubscriber, "MSG ");
    }

    public async ValueTask DisposeAsync()
    {
        _spokeSubscriber?.Dispose();
        _hubPublisher?.Dispose();
        await _hubCts.CancelAsync();
        await _spokeCts.CancelAsync();
        _hub.Dispose();
        _spoke.Dispose();
        _hubCts.Dispose();
        _spokeCts.Dispose();
    }

    private static async Task<string> ReadLineAsync(Socket sock)
    {
        var buf = new byte[4096];
        var n = await sock.ReceiveAsync(buf, SocketFlags.None);
        return Encoding.ASCII.GetString(buf, 0, n);
    }


}