natsdotnet/tests/NATS.Server.Tests/ClientPubSubTests.cs

// Go reference: golang/nats-server/server/client_test.go
// TestClientSimplePubSub (line 666), TestClientPubSubNoEcho (line 691),
// TestClientSimplePubSubWithReply (line 712), TestClientNoBodyPubSubWithReply (line 740),
// TestClientPubWithQueueSub (line 768)

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

namespace NATS.Server.Tests;

public class ClientPubSubTests : IAsyncLifetime
{
    private readonly NatsServer _server;
    private readonly int _port;
    private readonly CancellationTokenSource _cts = new();

    public ClientPubSubTests()
    {
        _port = GetFreePort();
        _server = new NatsServer(new NatsOptions { Port = _port }, NullLoggerFactory.Instance);
    }

    public async Task InitializeAsync()
    {
        _ = _server.StartAsync(_cts.Token);
        await _server.WaitForReadyAsync();
    }

    public async Task DisposeAsync()
    {
        await _cts.CancelAsync();
        _server.Dispose();
    }

    private static int GetFreePort()
    {
        using var sock = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        sock.Bind(new IPEndPoint(IPAddress.Loopback, 0));
        return ((IPEndPoint)sock.LocalEndPoint!).Port;
    }

    private async Task<Socket> ConnectClientAsync()
    {
        var sock = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await sock.ConnectAsync(IPAddress.Loopback, _port);
        return sock;
    }

    /// <summary>
    /// Reads from a socket until the accumulated data contains the expected substring.
    /// </summary>
    private static async Task<string> ReadUntilAsync(Socket sock, string expected, int timeoutMs = 5000)
    {
        using var cts = new CancellationTokenSource(timeoutMs);
        var sb = new StringBuilder();
        var buf = new byte[4096];
        while (!sb.ToString().Contains(expected))
        {
            var n = await sock.ReceiveAsync(buf, SocketFlags.None, cts.Token);
            if (n == 0) break;
            sb.Append(Encoding.ASCII.GetString(buf, 0, n));
        }
        return sb.ToString();
    }

    // Go reference: TestClientSimplePubSub (client_test.go line 666)
    // SUB foo 1, PUB foo 5\r\nhello — subscriber receives MSG foo 1 5\r\nhello
    [Fact]
    public async Task Simple_pub_sub_delivers_message()
    {
        using var client = await ConnectClientAsync();

        // Read INFO
        var buf = new byte[4096];
        await client.ReceiveAsync(buf, SocketFlags.None);

        // CONNECT, SUB, PUB, then PING to flush delivery
        await client.SendAsync(Encoding.ASCII.GetBytes(
            "CONNECT {}\r\nSUB foo 1\r\nPUB foo 5\r\nhello\r\nPING\r\n"));

        // Read until we see the message payload (delivered before PONG)
        var response = await ReadUntilAsync(client, "hello\r\n");

        // MSG line: MSG foo 1 5\r\nhello\r\n
        response.ShouldContain("MSG foo 1 5\r\nhello\r\n");
    }

    // Go reference: TestClientPubSubNoEcho (client_test.go line 691)
    // CONNECT {"echo":false} — publishing client does NOT receive its own messages
    [Fact]
    public async Task Pub_sub_no_echo_suppresses_own_messages()
    {
        using var client = await ConnectClientAsync();

        // Read INFO
        var buf = new byte[4096];
        await client.ReceiveAsync(buf, SocketFlags.None);

        // Connect with echo=false, then SUB+PUB on same connection, then PING
        await client.SendAsync(Encoding.ASCII.GetBytes(
            "CONNECT {\"echo\":false}\r\nSUB foo 1\r\nPUB foo 5\r\nhello\r\nPING\r\n"));

        // With echo=false the server must not deliver the message back to the publisher.
        // The first line we receive should be PONG, not MSG.
        var response = await ReadUntilAsync(client, "PONG\r\n");

        response.ShouldStartWith("PONG\r\n");
        response.ShouldNotContain("MSG");
    }

    // Go reference: TestClientSimplePubSubWithReply (client_test.go line 712)
    // PUB foo bar 5\r\nhello — subscriber receives MSG foo 1 bar 5\r\nhello (reply subject included)
    [Fact]
    public async Task Pub_sub_with_reply_subject()
    {
        using var client = await ConnectClientAsync();

        // Read INFO
        var buf = new byte[4096];
        await client.ReceiveAsync(buf, SocketFlags.None);

        // PUB with reply subject "bar"
        await client.SendAsync(Encoding.ASCII.GetBytes(
            "CONNECT {}\r\nSUB foo 1\r\nPUB foo bar 5\r\nhello\r\nPING\r\n"));

        var response = await ReadUntilAsync(client, "hello\r\n");

        // MSG line must include the reply subject: MSG <subject> <sid> <reply> <#bytes>
        response.ShouldContain("MSG foo 1 bar 5\r\nhello\r\n");
    }

    // Go reference: TestClientNoBodyPubSubWithReply (client_test.go line 740)
    // PUB foo bar 0\r\n\r\n — zero-byte payload with reply subject
    [Fact]
    public async Task Empty_body_pub_sub_with_reply()
    {
        using var client = await ConnectClientAsync();

        // Read INFO
        var buf = new byte[4096];
        await client.ReceiveAsync(buf, SocketFlags.None);

        // PUB with reply subject and zero-length body
        await client.SendAsync(Encoding.ASCII.GetBytes(
            "CONNECT {}\r\nSUB foo 1\r\nPUB foo bar 0\r\n\r\nPING\r\n"));

        // Read until PONG — MSG should arrive before PONG
        var response = await ReadUntilAsync(client, "PONG\r\n");

        // MSG line: MSG foo 1 bar 0\r\n\r\n (empty body, still CRLF terminated)
        response.ShouldContain("MSG foo 1 bar 0\r\n");
    }

    // Go reference: TestClientPubWithQueueSub (client_test.go line 768)
    // Two queue subscribers in the same group on one connection — 100 publishes
    // distributed across both sids, each receiving at least 20 messages.
    [Fact]
    public async Task Queue_sub_distributes_messages()
    {
        const int num = 100;

        using var client = await ConnectClientAsync();

        // Read INFO
        var buf = new byte[4096];
        await client.ReceiveAsync(buf, SocketFlags.None);

        // CONNECT, two queue subs with different sids, PING to confirm
        await client.SendAsync(Encoding.ASCII.GetBytes(
            "CONNECT {}\r\nSUB foo g1 1\r\nSUB foo g1 2\r\nPING\r\n"));
        await ReadUntilAsync(client, "PONG\r\n");

        // Publish 100 messages, then PING to flush all deliveries
        var pubSb = new StringBuilder();
        for (int i = 0; i < num; i++)
            pubSb.Append("PUB foo 5\r\nhello\r\n");
        pubSb.Append("PING\r\n");
        await client.SendAsync(Encoding.ASCII.GetBytes(pubSb.ToString()));

        // Read until PONG — all MSGs arrive before the PONG
        var response = await ReadUntilAsync(client, "PONG\r\n");

        // Count deliveries per sid
        var n1 = Regex.Matches(response, @"MSG foo 1 5").Count;
        var n2 = Regex.Matches(response, @"MSG foo 2 5").Count;

        (n1 + n2).ShouldBe(num);
        n1.ShouldBeGreaterThanOrEqualTo(20);
        n2.ShouldBeGreaterThanOrEqualTo(20);
    }
}