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

using SystemWebSocket = System.Net.WebSockets.WebSocket;
using System.Net.WebSockets;
using NSubstitute;
using NATS.Server.LeafNodes;

namespace NATS.Server.Tests.LeafNodes;

/// <summary>
/// Unit tests for <see cref="WebSocketStreamAdapter"/> (Gap 12.5).
/// Verifies stream capability flags, read/write delegation to WebSocket,
/// telemetry counters, and IsConnected state reflection.
/// Go reference: leafnode.go wsCreateLeafConnection, client.go wsRead/wsWrite.
/// </summary>
public class LeafWebSocketTests
{
    // -------------------------------------------------------------------------
    // Helper
    // -------------------------------------------------------------------------

    private static SystemWebSocket CreateMockWebSocket(byte[]? readData = null)
    {
        var ws = Substitute.For<SystemWebSocket>();
        ws.State.Returns(WebSocketState.Open);

        if (readData != null)
        {
            ws.ReceiveAsync(Arg.Any<Memory<byte>>(), Arg.Any<CancellationToken>())
                .Returns(callInfo =>
                {
                    var mem = callInfo.ArgAt<Memory<byte>>(0);
                    var toCopy = Math.Min(readData.Length, mem.Length);
                    readData.AsSpan(0, toCopy).CopyTo(mem.Span);
                    return new ValueTask<ValueWebSocketReceiveResult>(
                        new ValueWebSocketReceiveResult(toCopy, WebSocketMessageType.Binary, true));
                });
        }

        return ws;
    }

    // -------------------------------------------------------------------------
    // Tests 1-3: Stream capability flags
    // -------------------------------------------------------------------------

    // Go reference: client.go wsRead — reads are supported
    [Fact]
    public void CanRead_ReturnsTrue()
    {
        var ws = CreateMockWebSocket();
        var adapter = new WebSocketStreamAdapter(ws);

        adapter.CanRead.ShouldBeTrue();
    }

    // Go reference: client.go wsWrite — writes are supported
    [Fact]
    public void CanWrite_ReturnsTrue()
    {
        var ws = CreateMockWebSocket();
        var adapter = new WebSocketStreamAdapter(ws);

        adapter.CanWrite.ShouldBeTrue();
    }

    // Go reference: leafnode.go wsCreateLeafConnection — WebSocket is not seekable
    [Fact]
    public void CanSeek_ReturnsFalse()
    {
        var ws = CreateMockWebSocket();
        var adapter = new WebSocketStreamAdapter(ws);

        adapter.CanSeek.ShouldBeFalse();
    }

    // -------------------------------------------------------------------------
    // Test 4: ReadAsync delegates to WebSocket
    // -------------------------------------------------------------------------

    // Go reference: client.go wsRead — receive next message from WebSocket
    [Fact]
    public async Task ReadAsync_ReceivesFromWebSocket()
    {
        var expected = "hello"u8.ToArray();
        var ws = CreateMockWebSocket(readData: expected);
        var adapter = new WebSocketStreamAdapter(ws);

        var buffer = new byte[16];
        var read = await adapter.ReadAsync(buffer, 0, buffer.Length, CancellationToken.None);

        read.ShouldBe(expected.Length);
        buffer[..read].ShouldBe(expected);
    }

    // -------------------------------------------------------------------------
    // Test 5: WriteAsync delegates to WebSocket
    // -------------------------------------------------------------------------

    // Go reference: client.go wsWrite — send data as a single binary frame
    [Fact]
    public async Task WriteAsync_SendsToWebSocket()
    {
        var ws = Substitute.For<SystemWebSocket>();
        ws.State.Returns(WebSocketState.Open);
        var capturedData = new List<byte>();

        ws.SendAsync(
                Arg.Any<ReadOnlyMemory<byte>>(),
                WebSocketMessageType.Binary,
                true,
                Arg.Any<CancellationToken>())
            .Returns(callInfo =>
            {
                var mem = callInfo.ArgAt<ReadOnlyMemory<byte>>(0);
                capturedData.AddRange(mem.ToArray());
                return ValueTask.CompletedTask;
            });

        var adapter = new WebSocketStreamAdapter(ws);
        var payload = "world"u8.ToArray();

        await adapter.WriteAsync(payload, 0, payload.Length, CancellationToken.None);

        await ws.Received(1).SendAsync(
            Arg.Any<ReadOnlyMemory<byte>>(),
            WebSocketMessageType.Binary,
            true,
            Arg.Any<CancellationToken>());
        capturedData.ShouldBe(payload);
    }

    // -------------------------------------------------------------------------
    // Test 6: BytesRead tracking
    // -------------------------------------------------------------------------

    // Go reference: client.go wsRead — track inbound byte count
    [Fact]
    public async Task BytesRead_TracksTotal()
    {
        var payload = new byte[] { 1, 2, 3, 4, 5 };
        var ws = CreateMockWebSocket(readData: payload);
        var adapter = new WebSocketStreamAdapter(ws);

        var buffer = new byte[16];
        var bytesRead = await adapter.ReadAsync(buffer.AsMemory(), CancellationToken.None);

        bytesRead.ShouldBeGreaterThan(0);
        adapter.BytesRead.ShouldBe(payload.Length);
    }

    // -------------------------------------------------------------------------
    // Test 7: BytesWritten tracking
    // -------------------------------------------------------------------------

    // Go reference: client.go wsWrite — track outbound byte count
    [Fact]
    public async Task BytesWritten_TracksTotal()
    {
        var ws = Substitute.For<SystemWebSocket>();
        ws.State.Returns(WebSocketState.Open);
        ws.SendAsync(Arg.Any<ReadOnlyMemory<byte>>(), Arg.Any<WebSocketMessageType>(), Arg.Any<bool>(), Arg.Any<CancellationToken>())
            .Returns(ValueTask.CompletedTask);

        var adapter = new WebSocketStreamAdapter(ws);
        var payload = new byte[] { 10, 20, 30 };

        await adapter.WriteAsync(payload, 0, payload.Length, CancellationToken.None);

        adapter.BytesWritten.ShouldBe(payload.Length);
    }

    // -------------------------------------------------------------------------
    // Test 8: MessagesRead counter
    // -------------------------------------------------------------------------

    // Go reference: client.go wsRead — each completed WebSocket message increments counter
    [Fact]
    public async Task MessagesRead_Incremented()
    {
        var payload = new byte[] { 0xAA, 0xBB };
        var ws = CreateMockWebSocket(readData: payload);
        var adapter = new WebSocketStreamAdapter(ws);

        var buffer = new byte[16];
        _ = await adapter.ReadAsync(buffer.AsMemory(), CancellationToken.None);

        adapter.MessagesRead.ShouldBe(1);
    }

    // -------------------------------------------------------------------------
    // Test 9: MessagesWritten counter
    // -------------------------------------------------------------------------

    // Go reference: client.go wsWrite — each SendAsync call is one message
    [Fact]
    public async Task MessagesWritten_Incremented()
    {
        var ws = Substitute.For<SystemWebSocket>();
        ws.State.Returns(WebSocketState.Open);
        ws.SendAsync(Arg.Any<ReadOnlyMemory<byte>>(), Arg.Any<WebSocketMessageType>(), Arg.Any<bool>(), Arg.Any<CancellationToken>())
            .Returns(ValueTask.CompletedTask);

        var adapter = new WebSocketStreamAdapter(ws);

        await adapter.WriteAsync(ReadOnlyMemory<byte>.Empty, CancellationToken.None);
        await adapter.WriteAsync(ReadOnlyMemory<byte>.Empty, CancellationToken.None);

        adapter.MessagesWritten.ShouldBe(2);
    }

    // -------------------------------------------------------------------------
    // Test 10: IsConnected reflects WebSocket.State
    // -------------------------------------------------------------------------

    // Go reference: leafnode.go wsCreateLeafConnection — connection liveness check
    [Fact]
    public void IsConnected_ReflectsWebSocketState()
    {
        var openWs = Substitute.For<SystemWebSocket>();
        openWs.State.Returns(WebSocketState.Open);

        var closedWs = Substitute.For<SystemWebSocket>();
        closedWs.State.Returns(WebSocketState.Closed);

        var openAdapter = new WebSocketStreamAdapter(openWs);
        var closedAdapter = new WebSocketStreamAdapter(closedWs);

        openAdapter.IsConnected.ShouldBeTrue();
        closedAdapter.IsConnected.ShouldBeFalse();
    }
}