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

using System.Net;
using System.Net.Sockets;
using System.Text;
using NATS.Server.LeafNodes;
using NATS.Server.Subscriptions;

namespace NATS.Server.LeafNodes.Tests.LeafNodes;

public class LeafConnectionParityBatch4Tests
{
    [Fact]
    public async Task SendLsPlus_with_queue_weight_writes_weighted_frame()
    {
        using var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        var endpoint = (IPEndPoint)listener.LocalEndpoint;

        using var remoteSocket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await remoteSocket.ConnectAsync(IPAddress.Loopback, endpoint.Port);
        using var leafSocket = await listener.AcceptSocketAsync();
        await using var leaf = new LeafConnection(leafSocket);

        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        var handshakeTask = leaf.PerformOutboundHandshakeAsync("LOCAL", timeout.Token);
        (await ReadLineAsync(remoteSocket, timeout.Token)).ShouldBe("LEAF LOCAL");
        await WriteLineAsync(remoteSocket, "LEAF REMOTE", timeout.Token);
        await handshakeTask;

        await leaf.SendLsPlusAsync("$G", "jobs.>", "workers", queueWeight: 3, timeout.Token);
        (await ReadLineAsync(remoteSocket, timeout.Token)).ShouldBe("LS+ $G jobs.> workers 3");
    }

    [Fact]
    public async Task ReadLoop_parses_queue_weight_from_ls_plus()
    {
        using var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        var endpoint = (IPEndPoint)listener.LocalEndpoint;

        using var remoteSocket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await remoteSocket.ConnectAsync(IPAddress.Loopback, endpoint.Port);
        using var leafSocket = await listener.AcceptSocketAsync();
        await using var leaf = new LeafConnection(leafSocket);

        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        var handshakeTask = leaf.PerformOutboundHandshakeAsync("LOCAL", timeout.Token);
        (await ReadLineAsync(remoteSocket, timeout.Token)).ShouldBe("LEAF LOCAL");
        await WriteLineAsync(remoteSocket, "LEAF REMOTE", timeout.Token);
        await handshakeTask;

        var received = new List<RemoteSubscription>();
        leaf.RemoteSubscriptionReceived = sub =>
        {
            received.Add(sub);
            return Task.CompletedTask;
        };
        leaf.StartLoop(timeout.Token);

        await WriteLineAsync(remoteSocket, "LS+ $G jobs.> workers 7", timeout.Token);
        await WaitForAsync(() => received.Count >= 1, timeout.Token);

        received[0].Subject.ShouldBe("jobs.>");
        received[0].Queue.ShouldBe("workers");
        received[0].QueueWeight.ShouldBe(7);
    }

    [Fact]
    public async Task ReadLoop_defaults_invalid_queue_weight_to_one()
    {
        using var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        var endpoint = (IPEndPoint)listener.LocalEndpoint;

        using var remoteSocket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await remoteSocket.ConnectAsync(IPAddress.Loopback, endpoint.Port);
        using var leafSocket = await listener.AcceptSocketAsync();
        await using var leaf = new LeafConnection(leafSocket);

        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        var handshakeTask = leaf.PerformOutboundHandshakeAsync("LOCAL", timeout.Token);
        (await ReadLineAsync(remoteSocket, timeout.Token)).ShouldBe("LEAF LOCAL");
        await WriteLineAsync(remoteSocket, "LEAF REMOTE", timeout.Token);
        await handshakeTask;

        var received = new List<RemoteSubscription>();
        leaf.RemoteSubscriptionReceived = sub =>
        {
            received.Add(sub);
            return Task.CompletedTask;
        };
        leaf.StartLoop(timeout.Token);

        await WriteLineAsync(remoteSocket, "LS+ $G jobs.> workers 0", timeout.Token);
        await WaitForAsync(() => received.Count >= 1, timeout.Token);

        received[0].QueueWeight.ShouldBe(1);
    }

    private static async Task<string> ReadLineAsync(Socket socket, CancellationToken ct)
    {
        var bytes = new List<byte>(64);
        var single = new byte[1];
        while (true)
        {
            var read = await socket.ReceiveAsync(single, SocketFlags.None, ct);
            if (read == 0)
                throw new IOException("Connection closed while reading line");
            if (single[0] == (byte)'\n')
                break;
            if (single[0] != (byte)'\r')
                bytes.Add(single[0]);
        }

        return Encoding.ASCII.GetString([.. bytes]);
    }

    private static Task WriteLineAsync(Socket socket, string line, CancellationToken ct)
        => socket.SendAsync(Encoding.ASCII.GetBytes($"{line}\r\n"), SocketFlags.None, ct).AsTask();

    private static async Task WaitForAsync(Func<bool> predicate, CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            if (predicate())
                return;

            await Task.Yield();
        }

        throw new TimeoutException("Timed out waiting for condition.");
    }
}