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

using NATS.Client.Core;
using NATS.E2E.Tests.Infrastructure;

namespace NATS.E2E.Tests;

public class ShutdownDrainTests
{
    [Fact]
    public async Task ClientDrain_CompletesInFlightMessages()
    {
        await using var server = new NatsServerProcess();
        await server.StartAsync();

        await using var pub = new NatsConnection(new NatsOpts { Url = $"nats://127.0.0.1:{server.Port}" });
        var sub = new NatsConnection(new NatsOpts { Url = $"nats://127.0.0.1:{server.Port}" });
        await using var _ = sub;

        await pub.ConnectAsync();
        await sub.ConnectAsync();

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

        for (var i = 0; i < 10; i++)
            await pub.PublishAsync($"e2e.drain.{i}", $"msg{i}");
        await pub.PingAsync();

        var received = new List<string?>();
        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));

        for (var i = 0; i < 10; i++)
        {
            var msg = await subscription.Msgs.ReadAsync(cts.Token);
            received.Add(msg.Data);
        }

        received.Count.ShouldBe(10);

        await sub.DisposeAsync();
    }

    [Fact]
    public async Task ServerShutdown_ClientDetectsDisconnection()
    {
        var server = new NatsServerProcess();
        await server.StartAsync();

        await using var client = new NatsConnection(new NatsOpts
        {
            Url = $"nats://127.0.0.1:{server.Port}",
            MaxReconnectRetry = 0,
        });
        await client.ConnectAsync();
        await client.PingAsync();

        client.ConnectionState.ShouldBe(NatsConnectionState.Open);

        var disconnected = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        client.ConnectionDisconnected += (_, _) =>
        {
            disconnected.TrySetResult();
            return ValueTask.CompletedTask;
        };

        await server.DisposeAsync();

        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        await disconnected.Task.WaitAsync(cts.Token);

        client.ConnectionState.ShouldNotBe(NatsConnectionState.Open);
    }
}