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

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

namespace NATS.Server.Tests;

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

    public ServerTests()
    {
        // Use random port
        _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;
    }

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

    /// <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();
    }

    [Fact]
    public async Task Server_accepts_connection_and_sends_INFO()
    {
        using var client = await ConnectClientAsync();
        var response = await ReadLineAsync(client);

        response.ShouldStartWith("INFO ");
    }

    [Fact]
    public async Task Server_basic_pubsub()
    {
        using var pub = await ConnectClientAsync();
        using var sub = await ConnectClientAsync();

        // Read INFO from both
        await ReadLineAsync(pub);
        await ReadLineAsync(sub);

        // CONNECT + SUB on subscriber, then PING to flush
        await sub.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nSUB foo 1\r\nPING\r\n"));
        var pong = await ReadLineAsync(sub);
        pong.ShouldContain("PONG");

        // CONNECT + PUB on publisher
        await pub.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPUB foo 5\r\nHello\r\n"));

        // Read MSG from subscriber (may arrive across multiple TCP segments)
        var msg = await ReadUntilAsync(sub, "Hello\r\n");
        msg.ShouldContain("MSG foo 1 5\r\nHello\r\n");
    }

    [Fact]
    public async Task Server_wildcard_matching()
    {
        using var pub = await ConnectClientAsync();
        using var sub = await ConnectClientAsync();

        await ReadLineAsync(pub);
        await ReadLineAsync(sub);

        // CONNECT + SUB on subscriber, then PING to flush
        await sub.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nSUB foo.* 1\r\nPING\r\n"));
        var pong = await ReadLineAsync(sub);
        pong.ShouldContain("PONG");

        await pub.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPUB foo.bar 5\r\nHello\r\n"));

        var buf = new byte[4096];
        var n = await sub.ReceiveAsync(buf, SocketFlags.None);
        var msg = Encoding.ASCII.GetString(buf, 0, n);

        msg.ShouldContain("MSG foo.bar 1 5\r\n");
    }

    [Fact]
    public async Task Server_pedantic_rejects_invalid_publish_subject()
    {
        using var pub = await ConnectClientAsync();
        using var sub = await ConnectClientAsync();

        // Read INFO from both
        await ReadLineAsync(pub);
        await ReadLineAsync(sub);

        // Connect with pedantic mode ON
        await pub.SendAsync(Encoding.ASCII.GetBytes(
            "CONNECT {\"pedantic\":true}\r\nPING\r\n"));
        var pong = await ReadUntilAsync(pub, "PONG");

        // Subscribe on sub
        await sub.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nSUB foo.* 1\r\nPING\r\n"));
        await ReadUntilAsync(sub, "PONG");

        // PUB with wildcard subject (invalid for publish)
        await pub.SendAsync(Encoding.ASCII.GetBytes("PUB foo.* 5\r\nHello\r\n"));

        // Publisher should get -ERR
        var errResponse = await ReadUntilAsync(pub, "-ERR", timeoutMs: 3000);
        errResponse.ShouldContain("-ERR 'Invalid Publish Subject'");
    }

    [Fact]
    public async Task Server_nonpedantic_allows_wildcard_publish_subject()
    {
        using var pub = await ConnectClientAsync();
        using var sub = await ConnectClientAsync();

        await ReadLineAsync(pub);
        await ReadLineAsync(sub);

        // Connect without pedantic mode (default)
        await sub.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nSUB foo.* 1\r\nPING\r\n"));
        await ReadUntilAsync(sub, "PONG");

        await pub.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPUB foo.* 5\r\nHello\r\n"));

        // Sub should still receive the message (no validation in non-pedantic mode)
        var msg = await ReadUntilAsync(sub, "Hello\r\n");
        msg.ShouldContain("MSG foo.* 1 5\r\nHello\r\n");
    }

    [Fact]
    public async Task Server_rejects_max_payload_violation()
    {
        // Create server with tiny max payload
        var port = GetFreePort();
        using var cts = new CancellationTokenSource();
        var server = new NatsServer(new NatsOptions { Port = port, MaxPayload = 10 }, NullLoggerFactory.Instance);
        _ = server.StartAsync(cts.Token);
        await server.WaitForReadyAsync();

        try
        {
            var client = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
            await client.ConnectAsync(IPAddress.Loopback, port);

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

            await client.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\n"));

            // Send PUB with payload larger than MaxPayload (10 bytes)
            await client.SendAsync(Encoding.ASCII.GetBytes("PUB foo 20\r\n12345678901234567890\r\n"));

            using var readCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
            var n = await client.ReceiveAsync(buf, SocketFlags.None, readCts.Token);
            var response = Encoding.ASCII.GetString(buf, 0, n);
            response.ShouldContain("-ERR 'Maximum Payload Violation'");

            // Connection should be closed
            n = await client.ReceiveAsync(buf, SocketFlags.None, readCts.Token);
            n.ShouldBe(0);

            client.Dispose();
        }
        finally
        {
            await cts.CancelAsync();
            server.Dispose();
        }
    }
}

public class EphemeralPortTests
{
    [Fact]
    public async Task Server_resolves_ephemeral_port()
    {
        using var cts = new CancellationTokenSource();
        var server = new NatsServer(new NatsOptions { Port = 0 }, NullLoggerFactory.Instance);

        _ = server.StartAsync(cts.Token);
        await server.WaitForReadyAsync();

        try
        {
            // Port should have been resolved to a real port
            server.Port.ShouldBeGreaterThan(0);

            // Connect a raw socket to prove the port actually works
            using var client = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
            await client.ConnectAsync(IPAddress.Loopback, server.Port);

            var buf = new byte[4096];
            var n = await client.ReceiveAsync(buf, SocketFlags.None);
            var response = Encoding.ASCII.GetString(buf, 0, n);
            response.ShouldStartWith("INFO ");
        }
        finally
        {
            await cts.CancelAsync();
            server.Dispose();
        }
    }
}

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

    public MaxConnectionsTests()
    {
        _port = GetFreePort();
        _server = new NatsServer(new NatsOptions { Port = _port, MaxConnections = 2 }, 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;
    }

    [Fact]
    public async Task Server_rejects_connection_when_max_reached()
    {
        using var readCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));

        // Connect two clients (at limit)
        var client1 = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await client1.ConnectAsync(IPAddress.Loopback, _port);
        var buf = new byte[4096];
        var n = await client1.ReceiveAsync(buf, SocketFlags.None, readCts.Token);
        Encoding.ASCII.GetString(buf, 0, n).ShouldStartWith("INFO ");

        var client2 = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await client2.ConnectAsync(IPAddress.Loopback, _port);
        n = await client2.ReceiveAsync(buf, SocketFlags.None, readCts.Token);
        Encoding.ASCII.GetString(buf, 0, n).ShouldStartWith("INFO ");

        // Third client should be rejected
        var client3 = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await client3.ConnectAsync(IPAddress.Loopback, _port);

        n = await client3.ReceiveAsync(buf, SocketFlags.None, readCts.Token);
        var response = Encoding.ASCII.GetString(buf, 0, n);
        response.ShouldContain("-ERR 'maximum connections exceeded'");

        // Connection should be closed
        n = await client3.ReceiveAsync(buf, SocketFlags.None, readCts.Token);
        n.ShouldBe(0);

        client1.Dispose();
        client2.Dispose();
        client3.Dispose();
    }
}

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

    public PingKeepaliveTests()
    {
        _port = GetFreePort();
        // Short intervals for testing: 500ms ping interval, 2 max pings out
        _server = new NatsServer(
            new NatsOptions
            {
                Port = _port,
                PingInterval = TimeSpan.FromMilliseconds(500),
                MaxPingsOut = 2,
            },
            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 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();
    }

    [Fact]
    public async Task Server_sends_PING_after_inactivity()
    {
        var client = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await client.ConnectAsync(IPAddress.Loopback, _port);

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

        // Send CONNECT to start keepalive
        await client.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\n"));

        // Wait for server to send PING (should come within ~500ms)
        var response = await ReadUntilAsync(client, "PING", timeoutMs: 3000);
        response.ShouldContain("PING");

        client.Dispose();
    }

    [Fact]
    public async Task Server_pong_resets_ping_counter()
    {
        var client = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await client.ConnectAsync(IPAddress.Loopback, _port);

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

        await client.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\n"));

        // Wait for first PING
        var response = await ReadUntilAsync(client, "PING", timeoutMs: 3000);
        response.ShouldContain("PING");

        // Respond with PONG — this resets the counter
        await client.SendAsync(Encoding.ASCII.GetBytes("PONG\r\n"));

        // Wait for next PING (counter reset, so we should get another one)
        response = await ReadUntilAsync(client, "PING", timeoutMs: 3000);
        response.ShouldContain("PING");

        // Respond again to keep alive
        await client.SendAsync(Encoding.ASCII.GetBytes("PONG\r\n"));

        // Client should still be alive — send a PING and expect PONG back
        await client.SendAsync(Encoding.ASCII.GetBytes("PING\r\n"));
        response = await ReadUntilAsync(client, "PONG", timeoutMs: 3000);
        response.ShouldContain("PONG");

        client.Dispose();
    }

    [Fact]
    public async Task Server_disconnects_stale_client()
    {
        var client = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await client.ConnectAsync(IPAddress.Loopback, _port);

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

        await client.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\n"));

        // Don't respond to PINGs — wait for stale disconnect
        // With 500ms interval and MaxPingsOut=2:
        // t=500ms: PING #1, pingsOut=1
        // t=1000ms: PING #2, pingsOut=2
        // t=1500ms: pingsOut+1 > MaxPingsOut → -ERR 'Stale Connection' + close
        var sb = new StringBuilder();
        try
        {
            using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
            while (true)
            {
                var n = await client.ReceiveAsync(buf, SocketFlags.None, timeout.Token);
                if (n == 0) break;
                sb.Append(Encoding.ASCII.GetString(buf, 0, n));
            }
        }
        catch (OperationCanceledException)
        {
            // Timeout is acceptable — check what we got
        }

        var allData = sb.ToString();
        allData.ShouldContain("-ERR 'Stale Connection'");

        client.Dispose();
    }
}

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

    public CloseReasonTests()
    {
        _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;
    }

    [Fact]
    public async Task Client_close_reason_set_on_normal_disconnect()
    {
        // Connect a raw TCP client
        using var client = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await client.ConnectAsync(IPAddress.Loopback, _port);

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

        // Send CONNECT + PING, wait for PONG
        await client.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPING\r\n"));
        var sb = new StringBuilder();
        using var readCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        while (!sb.ToString().Contains("PONG"))
        {
            var n = await client.ReceiveAsync(buf, SocketFlags.None, readCts.Token);
            if (n == 0) break;
            sb.Append(Encoding.ASCII.GetString(buf, 0, n));
        }
        sb.ToString().ShouldContain("PONG");

        // Get the NatsClient from the server
        var natsClient = _server.GetClients().First();

        // Close the TCP socket (normal client disconnect)
        client.Shutdown(SocketShutdown.Both);
        client.Close();

        // Wait for the server to detect the disconnect
        await Task.Delay(500);

        // The close reason should be ClientClosed (normal disconnect falls through to finally)
        natsClient.CloseReason.ShouldBe(ClientClosedReason.ClientClosed);
    }
}

public class ServerIdentityTests
{
    [Fact]
    public void Server_creates_system_account()
    {
        var server = new NatsServer(new NatsOptions { Port = 0 }, NullLoggerFactory.Instance);
        server.SystemAccount.ShouldNotBeNull();
        server.SystemAccount.Name.ShouldBe("$SYS");
        server.Dispose();
    }

    [Fact]
    public void Server_generates_nkey_identity()
    {
        var server = new NatsServer(new NatsOptions { Port = 0 }, NullLoggerFactory.Instance);
        server.ServerNKey.ShouldNotBeNullOrEmpty();
        // Server NKey public keys start with 'N'
        server.ServerNKey[0].ShouldBe('N');
        server.Dispose();
    }
}

public class FlushBeforeCloseTests
{
    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 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();
    }

    [Fact]
    public async Task Shutdown_flushes_pending_data_to_clients()
    {
        var port = GetFreePort();
        var server = new NatsServer(new NatsOptions { Port = port }, NullLoggerFactory.Instance);
        _ = server.StartAsync(CancellationToken.None);
        await server.WaitForReadyAsync();

        try
        {
            // Connect a subscriber via raw socket
            using var sub = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
            await sub.ConnectAsync(IPAddress.Loopback, port);

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

            // Subscribe to "foo"
            await sub.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nSUB foo 1\r\nPING\r\n"));
            var pong = await ReadUntilAsync(sub, "PONG");
            pong.ShouldContain("PONG");

            // Connect a publisher
            using var pub = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
            await pub.ConnectAsync(IPAddress.Loopback, port);
            await pub.ReceiveAsync(buf, SocketFlags.None); // INFO

            // Publish "Hello" to "foo"
            await pub.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPUB foo 5\r\nHello\r\n"));

            // Wait briefly for delivery
            await Task.Delay(200);

            // Read from subscriber to verify MSG was received
            var msg = await ReadUntilAsync(sub, "Hello\r\n");
            msg.ShouldContain("MSG foo 1 5\r\nHello\r\n");
        }
        finally
        {
            await server.ShutdownAsync();
            server.Dispose();
        }
    }
}

public class GracefulShutdownTests
{
    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;
    }

    [Fact]
    public async Task ShutdownAsync_disconnects_all_clients()
    {
        var port = GetFreePort();
        var server = new NatsServer(new NatsOptions { Port = port }, NullLoggerFactory.Instance);
        _ = server.StartAsync(CancellationToken.None);
        await server.WaitForReadyAsync();

        // Connect 2 raw TCP clients
        using var client1 = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await client1.ConnectAsync(IPAddress.Loopback, port);
        var buf = new byte[4096];
        await client1.ReceiveAsync(buf, SocketFlags.None); // INFO

        using var client2 = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await client2.ConnectAsync(IPAddress.Loopback, port);
        await client2.ReceiveAsync(buf, SocketFlags.None); // INFO

        // Send CONNECT so both are registered
        await client1.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPING\r\n"));
        await client2.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPING\r\n"));

        // Wait for PONG from both (confirming they are registered)
        using var readCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        await client1.ReceiveAsync(buf, SocketFlags.None, readCts.Token);
        await client2.ReceiveAsync(buf, SocketFlags.None, readCts.Token);

        server.ClientCount.ShouldBe(2);

        await server.ShutdownAsync();

        server.ClientCount.ShouldBe(0);
        server.Dispose();
    }

    [Fact]
    public async Task WaitForShutdown_blocks_until_shutdown()
    {
        var port = GetFreePort();
        var server = new NatsServer(new NatsOptions { Port = port }, NullLoggerFactory.Instance);
        _ = server.StartAsync(CancellationToken.None);
        await server.WaitForReadyAsync();

        // Start WaitForShutdown in background
        var waitTask = Task.Run(() => server.WaitForShutdown());

        // Give it a moment -- it should NOT complete yet
        await Task.Delay(200);
        waitTask.IsCompleted.ShouldBeFalse();

        // Trigger shutdown
        await server.ShutdownAsync();

        // WaitForShutdown should complete within 5 seconds
        var completed = await Task.WhenAny(waitTask, Task.Delay(TimeSpan.FromSeconds(5)));
        completed.ShouldBe(waitTask);

        server.Dispose();
    }

    [Fact]
    public async Task ShutdownAsync_is_idempotent()
    {
        var port = GetFreePort();
        var server = new NatsServer(new NatsOptions { Port = port }, NullLoggerFactory.Instance);
        _ = server.StartAsync(CancellationToken.None);
        await server.WaitForReadyAsync();

        // Call ShutdownAsync 3 times -- should not throw
        await server.ShutdownAsync();
        await server.ShutdownAsync();
        await server.ShutdownAsync();

        server.IsShuttingDown.ShouldBeTrue();
        server.Dispose();
    }

    [Fact]
    public async Task Accept_loop_waits_for_active_clients()
    {
        var port = GetFreePort();
        var server = new NatsServer(new NatsOptions { Port = port }, NullLoggerFactory.Instance);
        _ = server.StartAsync(CancellationToken.None);
        await server.WaitForReadyAsync();

        // Connect a client
        using var client = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        await client.ConnectAsync(IPAddress.Loopback, port);
        var buf = new byte[4096];
        await client.ReceiveAsync(buf, SocketFlags.None); // INFO
        await client.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPING\r\n"));
        using var readCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        await client.ReceiveAsync(buf, SocketFlags.None, readCts.Token); // PONG

        // ShutdownAsync should complete within 10 seconds (doesn't hang)
        var shutdownTask = server.ShutdownAsync();
        var completed = await Task.WhenAny(shutdownTask, Task.Delay(TimeSpan.FromSeconds(10)));
        completed.ShouldBe(shutdownTask);

        server.Dispose();
    }
}

public class LameDuckTests
{
    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;
    }

    [Fact]
    public async Task LameDuckShutdown_stops_accepting_new_connections()
    {
        var port = GetFreePort();
        var server = new NatsServer(
            new NatsOptions
            {
                Port = port,
                LameDuckDuration = TimeSpan.FromSeconds(3),
                LameDuckGracePeriod = TimeSpan.FromMilliseconds(500),
            },
            NullLoggerFactory.Instance);

        _ = server.StartAsync(CancellationToken.None);
        await server.WaitForReadyAsync();

        try
        {
            // Connect 1 client
            using var client1 = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
            await client1.ConnectAsync(IPAddress.Loopback, port);
            var buf = new byte[4096];
            await client1.ReceiveAsync(buf, SocketFlags.None); // INFO
            await client1.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPING\r\n"));
            using var readCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
            await client1.ReceiveAsync(buf, SocketFlags.None, readCts.Token); // PONG

            // Start lame duck (don't await yet)
            var lameDuckTask = server.LameDuckShutdownAsync();

            // Wait briefly for listener to close
            await Task.Delay(300);

            // Verify lame duck mode is active
            server.IsLameDuckMode.ShouldBeTrue();

            // Try connecting a new client -- should fail (connection refused)
            using var client2 = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
            var connectAction = async () =>
            {
                await client2.ConnectAsync(IPAddress.Loopback, port);
            };
            await connectAction.ShouldThrowAsync<SocketException>();

            // Await the lame duck task with timeout
            var completed = await Task.WhenAny(lameDuckTask, Task.Delay(TimeSpan.FromSeconds(15)));
            completed.ShouldBe(lameDuckTask);
        }
        finally
        {
            server.Dispose();
        }
    }

    [Fact]
    public async Task LameDuckShutdown_eventually_closes_all_clients()
    {
        var port = GetFreePort();
        var server = new NatsServer(
            new NatsOptions
            {
                Port = port,
                LameDuckDuration = TimeSpan.FromSeconds(2),
                LameDuckGracePeriod = TimeSpan.FromMilliseconds(200),
            },
            NullLoggerFactory.Instance);

        _ = server.StartAsync(CancellationToken.None);
        await server.WaitForReadyAsync();

        try
        {
            // Connect 3 clients via raw sockets
            var clients = new List<Socket>();
            var buf = new byte[4096];
            for (int i = 0; i < 3; i++)
            {
                var sock = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
                await sock.ConnectAsync(IPAddress.Loopback, port);
                await sock.ReceiveAsync(buf, SocketFlags.None); // INFO
                await sock.SendAsync(Encoding.ASCII.GetBytes("CONNECT {}\r\nPING\r\n"));
                using var readCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
                await sock.ReceiveAsync(buf, SocketFlags.None, readCts.Token); // PONG
                clients.Add(sock);
            }

            server.ClientCount.ShouldBe(3);

            // Await LameDuckShutdownAsync
            var lameDuckTask = server.LameDuckShutdownAsync();
            var completed = await Task.WhenAny(lameDuckTask, Task.Delay(TimeSpan.FromSeconds(15)));
            completed.ShouldBe(lameDuckTask);

            server.ClientCount.ShouldBe(0);

            foreach (var sock in clients)
                sock.Dispose();
        }
        finally
        {
            server.Dispose();
        }
    }
}

public class PidFileTests : IDisposable
{
    private readonly string _tempDir = Path.Combine(Path.GetTempPath(), $"nats-test-{Guid.NewGuid():N}");

    public PidFileTests() => Directory.CreateDirectory(_tempDir);

    public void Dispose()
    {
        if (Directory.Exists(_tempDir))
            Directory.Delete(_tempDir, recursive: true);
    }

    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;
    }

    [Fact]
    public async Task Server_writes_pid_file_on_startup()
    {
        var pidFile = Path.Combine(_tempDir, "nats.pid");
        var port = GetFreePort();
        var server = new NatsServer(new NatsOptions { Port = port, PidFile = pidFile }, NullLoggerFactory.Instance);
        _ = server.StartAsync(CancellationToken.None);
        await server.WaitForReadyAsync();

        File.Exists(pidFile).ShouldBeTrue();
        var content = await File.ReadAllTextAsync(pidFile);
        int.Parse(content).ShouldBe(Environment.ProcessId);

        await server.ShutdownAsync();
        File.Exists(pidFile).ShouldBeFalse();

        server.Dispose();
    }

    [Fact]
    public async Task Server_writes_ports_file_on_startup()
    {
        var port = GetFreePort();
        var server = new NatsServer(new NatsOptions { Port = port, PortsFileDir = _tempDir }, NullLoggerFactory.Instance);
        _ = server.StartAsync(CancellationToken.None);
        await server.WaitForReadyAsync();

        var portsFiles = Directory.GetFiles(_tempDir, "*.ports");
        portsFiles.Length.ShouldBe(1);

        var content = await File.ReadAllTextAsync(portsFiles[0]);
        content.ShouldContain($"\"client\":{port}");

        await server.ShutdownAsync();
        Directory.GetFiles(_tempDir, "*.ports").Length.ShouldBe(0);

        server.Dispose();
    }
}