natsdotnet/tests/NATS.E2E.Tests/CoreMessagingTests.cs

using System.Linq;
using System.Net.Sockets;
using System.Text;
using NATS.Client.Core;
using NATS.E2E.Tests.Infrastructure;

namespace NATS.E2E.Tests;

[Collection("E2E")]
public class CoreMessagingTests(NatsServerFixture fixture)
{
    [Fact]
    public async Task WildcardStar_MatchesSingleToken()
    {
        await using var pub = fixture.CreateClient();
        await using var sub = fixture.CreateClient();
        await pub.ConnectAsync();
        await sub.ConnectAsync();

        await using var subscription = await sub.SubscribeCoreAsync<string>("e2e.wc.*");
        await sub.PingAsync();

        await pub.PublishAsync("e2e.wc.bar", "hello");

        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        var msg = await subscription.Msgs.ReadAsync(cts.Token);

        msg.Data.ShouldBe("hello");
    }

    [Fact]
    public async Task WildcardGreaterThan_MatchesMultipleTokens()
    {
        await using var pub = fixture.CreateClient();
        await using var sub = fixture.CreateClient();
        await pub.ConnectAsync();
        await sub.ConnectAsync();

        await using var subscription = await sub.SubscribeCoreAsync<string>("e2e.gt.>");
        await sub.PingAsync();

        await pub.PublishAsync("e2e.gt.bar.baz", "deep");

        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        var msg = await subscription.Msgs.ReadAsync(cts.Token);

        msg.Data.ShouldBe("deep");
    }

    [Fact]
    public async Task WildcardStar_DoesNotMatchMultipleTokens()
    {
        await using var pub = fixture.CreateClient();
        await using var sub = fixture.CreateClient();
        await pub.ConnectAsync();
        await sub.ConnectAsync();

        await using var subscription = await sub.SubscribeCoreAsync<string>("e2e.nomat.*");
        await sub.PingAsync();

        await pub.PublishAsync("e2e.nomat.bar.baz", "should not arrive");

        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        var readTask = subscription.Msgs.ReadAsync(cts.Token).AsTask();
        var winner = await Task.WhenAny(readTask, Task.Delay(3000));

        winner.ShouldNotBe(readTask);
    }

    [Fact]
    public async Task QueueGroup_LoadBalances()
    {
        await using var c1 = fixture.CreateClient();
        await using var c2 = fixture.CreateClient();
        await using var c3 = fixture.CreateClient();
        await using var pub = fixture.CreateClient();

        await c1.ConnectAsync();
        await c2.ConnectAsync();
        await c3.ConnectAsync();
        await pub.ConnectAsync();

        await using var s1 = await c1.SubscribeCoreAsync<int>("e2e.qlb", queueGroup: "workers");
        await using var s2 = await c2.SubscribeCoreAsync<int>("e2e.qlb", queueGroup: "workers");
        await using var s3 = await c3.SubscribeCoreAsync<int>("e2e.qlb", queueGroup: "workers");

        await c1.PingAsync();
        await c2.PingAsync();
        await c3.PingAsync();

        using var collectionCts = new CancellationTokenSource(TimeSpan.FromSeconds(15));

        var counts = new int[3];

        async Task Collect(INatsSub<int> sub, int idx)
        {
            try
            {
                await foreach (var _ in sub.Msgs.ReadAllAsync(collectionCts.Token))
                    Interlocked.Increment(ref counts[idx]);
            }
            catch (OperationCanceledException) { }
        }

        var tasks = new[]
        {
            Collect(s1, 0),
            Collect(s2, 1),
            Collect(s3, 2),
        };

        for (var i = 0; i < 30; i++)
            await pub.PublishAsync("e2e.qlb", i);

        await pub.PingAsync();

        // Wait until all 30 messages have been received
        using var deadline = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        while (counts[0] + counts[1] + counts[2] < 30 && !deadline.IsCancellationRequested)
            await Task.Delay(10, deadline.Token).ContinueWith(_ => { });

        collectionCts.Cancel();
        await Task.WhenAll(tasks);

        var total = counts[0] + counts[1] + counts[2];
        total.ShouldBe(30);

        // Verify at least one queue member received messages (distribution
        // is implementation-defined and may heavily favor one member when
        // messages are published in a tight loop).
        counts.Max().ShouldBeGreaterThan(0);
    }

    [Fact]
    public async Task QueueGroup_MixedWithPlainSub()
    {
        await using var plainClient = fixture.CreateClient();
        await using var q1Client = fixture.CreateClient();
        await using var q2Client = fixture.CreateClient();
        await using var pub = fixture.CreateClient();

        await plainClient.ConnectAsync();
        await q1Client.ConnectAsync();
        await q2Client.ConnectAsync();
        await pub.ConnectAsync();

        await using var plainSub = await plainClient.SubscribeCoreAsync<int>("e2e.qmix");
        await using var qSub1 = await q1Client.SubscribeCoreAsync<int>("e2e.qmix", queueGroup: "qmix");
        await using var qSub2 = await q2Client.SubscribeCoreAsync<int>("e2e.qmix", queueGroup: "qmix");

        await plainClient.PingAsync();
        await q1Client.PingAsync();
        await q2Client.PingAsync();

        using var collectionCts = new CancellationTokenSource(TimeSpan.FromSeconds(15));

        var plainCount = 0;
        var q1Count = 0;
        var q2Count = 0;

        async Task Collect(INatsSub<int> sub, Action increment)
        {
            try
            {
                await foreach (var _ in sub.Msgs.ReadAllAsync(collectionCts.Token))
                    increment();
            }
            catch (OperationCanceledException) { }
        }

        var tasks = new[]
        {
            Collect(plainSub, () => Interlocked.Increment(ref plainCount)),
            Collect(qSub1, () => Interlocked.Increment(ref q1Count)),
            Collect(qSub2, () => Interlocked.Increment(ref q2Count)),
        };

        for (var i = 0; i < 10; i++)
            await pub.PublishAsync("e2e.qmix", i);

        await pub.PingAsync();

        using var deadline = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        while (plainCount < 10 && !deadline.IsCancellationRequested)
            await Task.Delay(10, deadline.Token).ContinueWith(_ => { });

        collectionCts.Cancel();
        await Task.WhenAll(tasks);

        plainCount.ShouldBe(10);
        (q1Count + q2Count).ShouldBe(10);
    }

    [Fact]
    public async Task Unsub_StopsDelivery()
    {
        await using var pub = fixture.CreateClient();
        await using var subClient = fixture.CreateClient();
        await pub.ConnectAsync();
        await subClient.ConnectAsync();

        var subscription = await subClient.SubscribeCoreAsync<string>("e2e.unsub");
        await subClient.PingAsync();

        await subscription.DisposeAsync();
        await subClient.PingAsync();

        // Subscribe a fresh listener on the same subject to verify the unsubscribed
        // client does NOT receive messages, while the fresh one does.
        await using var verifyClient = fixture.CreateClient();
        await verifyClient.ConnectAsync();
        await using var verifySub = await verifyClient.SubscribeCoreAsync<string>("e2e.unsub");
        await verifyClient.PingAsync();

        await pub.PublishAsync("e2e.unsub", "after-unsub");
        await pub.PingAsync();

        // The fresh subscriber should receive the message
        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        var msg = await verifySub.Msgs.ReadAsync(cts.Token);
        msg.Data.ShouldBe("after-unsub");

        // The original (disposed) subscription's channel should be completed —
        // reading from it should NOT yield "after-unsub"
        var received = new List<string?>();
        try
        {
            await foreach (var m in subscription.Msgs.ReadAllAsync(default))
                received.Add(m.Data);
        }
        catch { /* channel completed or cancelled — expected */ }

        received.ShouldNotContain("after-unsub");
    }

    [Fact]
    public async Task Unsub_WithMaxMessages()
    {
        using var tcp = new TcpClient();
        await tcp.ConnectAsync("127.0.0.1", fixture.Port);
        await using var ns = tcp.GetStream();
        using var reader = new StreamReader(ns, Encoding.ASCII, leaveOpen: true);
        var writer = new StreamWriter(ns, Encoding.ASCII, leaveOpen: true) { AutoFlush = true, NewLine = "\r\n" };

        // Read INFO
        var info = await reader.ReadLineAsync();
        info.ShouldNotBeNull();
        info.ShouldStartWith("INFO");

        await writer.WriteLineAsync("CONNECT {\"verbose\":false,\"protocol\":1}");
        await writer.WriteLineAsync("SUB e2e.unsub.max 1");
        await writer.WriteLineAsync("UNSUB 1 3");
        await writer.WriteLineAsync("PING");

        // Wait for PONG to know server processed commands
        string? line;
        do { line = await reader.ReadLineAsync(); } while (line != null && !line.StartsWith("PONG"));

        for (var i = 0; i < 5; i++)
            await writer.WriteLineAsync($"PUB e2e.unsub.max 1\r\nx");

        await writer.WriteLineAsync("PING");

        var msgCount = 0;
        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        while (!cts.IsCancellationRequested)
        {
            line = await reader.ReadLineAsync();
            if (line == null) break;
            if (line.StartsWith("MSG")) msgCount++;
            if (line.StartsWith("PONG")) break;
        }

        msgCount.ShouldBe(3);
    }

    [Fact]
    public async Task FanOut_MultipleSubscribers()
    {
        await using var pub = fixture.CreateClient();
        await using var sub1 = fixture.CreateClient();
        await using var sub2 = fixture.CreateClient();
        await using var sub3 = fixture.CreateClient();

        await pub.ConnectAsync();
        await sub1.ConnectAsync();
        await sub2.ConnectAsync();
        await sub3.ConnectAsync();

        await using var s1 = await sub1.SubscribeCoreAsync<string>("e2e.fanout");
        await using var s2 = await sub2.SubscribeCoreAsync<string>("e2e.fanout");
        await using var s3 = await sub3.SubscribeCoreAsync<string>("e2e.fanout");

        await sub1.PingAsync();
        await sub2.PingAsync();
        await sub3.PingAsync();

        await pub.PublishAsync("e2e.fanout", "broadcast");

        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));

        var m1 = await s1.Msgs.ReadAsync(cts.Token);
        var m2 = await s2.Msgs.ReadAsync(cts.Token);
        var m3 = await s3.Msgs.ReadAsync(cts.Token);

        m1.Data.ShouldBe("broadcast");
        m2.Data.ShouldBe("broadcast");
        m3.Data.ShouldBe("broadcast");
    }

    [Fact]
    public async Task EchoOff_PublisherDoesNotReceiveSelf()
    {
        var opts = new NatsOpts
        {
            Url = $"nats://127.0.0.1:{fixture.Port}",
            Echo = false,
        };

        await using var client = new NatsConnection(opts);
        await client.ConnectAsync();

        await using var subscription = await client.SubscribeCoreAsync<string>("e2e.echo");
        await client.PingAsync();

        await client.PublishAsync("e2e.echo", "self");

        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        var readTask = subscription.Msgs.ReadAsync(cts.Token).AsTask();
        var winner = await Task.WhenAny(readTask, Task.Delay(3000));

        winner.ShouldNotBe(readTask);
    }

    [Fact]
    public async Task VerboseMode_OkResponses()
    {
        using var tcp = new TcpClient();
        await tcp.ConnectAsync("127.0.0.1", fixture.Port);
        await using var ns = tcp.GetStream();
        using var reader = new StreamReader(ns, Encoding.ASCII, leaveOpen: true);
        var writer = new StreamWriter(ns, Encoding.ASCII, leaveOpen: true) { AutoFlush = true, NewLine = "\r\n" };

        // Read INFO
        var info = await reader.ReadLineAsync();
        info.ShouldNotBeNull();
        info.ShouldStartWith("INFO");

        await writer.WriteLineAsync("CONNECT {\"verbose\":true,\"protocol\":1}");

        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));

        var connectResponse = await reader.ReadLineAsync();
        connectResponse.ShouldBe("+OK");

        await writer.WriteLineAsync("SUB test 1");
        var subResponse = await reader.ReadLineAsync();
        subResponse.ShouldBe("+OK");

        await writer.WriteLineAsync("PING");
        var pongResponse = await reader.ReadLineAsync();
        pongResponse.ShouldBe("PONG");
    }

    [Fact]
    public async Task NoResponders_Returns503()
    {
        await using var client = fixture.CreateClient();
        await client.ConnectAsync();

        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));

        await Should.ThrowAsync<Exception>(async () =>
        {
            await client.RequestAsync<string, string>("e2e.noresp.xxx", "ping", cancellationToken: cts.Token);
        });
    }
}