zcbot/windows-node/Zcbot.WindowsNode.Tests/NodeProtocolClientTests.cs

using System.Collections.Concurrent;
using System.Net.WebSockets;
using System.Text;

namespace Zcbot.WindowsNode.Tests;

public sealed class NodeProtocolClientTests
{
    [Fact]
    public async Task SendAsyncSerializesConcurrentWriters()
    {
        var socket = new TestWebSocket(sendDelay: TimeSpan.FromMilliseconds(15));
        await using var client = new NodeProtocolClient(socket);

        await Task.WhenAll(Enumerable.Range(0, 8).Select(index =>
            client.SendAsync("job_state", new { index }, CancellationToken.None)));

        Assert.Equal(1, socket.MaximumConcurrentSends);
        Assert.Equal(8, socket.SentMessages.Count);
    }

    [Fact]
    public async Task ReceiveAsyncReassemblesTextFragments()
    {
        var socket = new TestWebSocket();
        socket.EnqueueText("{\"type\":\"job_", endOfMessage: false);
        socket.EnqueueText("cancel\",\"payload\":{\"job_id\":\"abc\"}}", endOfMessage: true);
        await using var client = new NodeProtocolClient(socket);

        var message = await client.ReceiveAsync(CancellationToken.None);

        Assert.NotNull(message);
        Assert.Equal("job_cancel", message.Type);
        Assert.Equal("abc", message.Payload.GetProperty("job_id").GetString());
    }

    [Fact]
    public async Task ReceiveAsyncTurnsIdentityRejectionIntoConfigurationFailure()
    {
        var socket = new TestWebSocket();
        socket.EnqueueClose((WebSocketCloseStatus)4003, "disabled");
        await using var client = new NodeProtocolClient(socket);

        var exception = await Assert.ThrowsAsync<NodeConfigurationException>(
            () => client.ReceiveAsync(CancellationToken.None));

        Assert.Contains("节点身份已被服务端拒绝", exception.Message);
        Assert.Contains("disabled", exception.Message);
    }

    private sealed class TestWebSocket : WebSocket
    {
        private readonly ConcurrentQueue<ReceiveFrame> receiveFrames = new();
        private readonly TimeSpan sendDelay;
        private int activeSends;
        private int maximumConcurrentSends;
        private WebSocketCloseStatus? closeStatus;
        private string? closeStatusDescription;
        private WebSocketState state = WebSocketState.Open;

        internal TestWebSocket(TimeSpan sendDelay = default)
        {
            this.sendDelay = sendDelay;
        }

        internal ConcurrentQueue<byte[]> SentMessages { get; } = new();

        internal int MaximumConcurrentSends => maximumConcurrentSends;

        public override WebSocketCloseStatus? CloseStatus => closeStatus;

        public override string? CloseStatusDescription => closeStatusDescription;

        public override WebSocketState State => state;

        public override string? SubProtocol => null;

        internal void EnqueueText(string value, bool endOfMessage)
        {
            receiveFrames.Enqueue(new ReceiveFrame(
                Encoding.UTF8.GetBytes(value),
                WebSocketMessageType.Text,
                endOfMessage,
                null,
                null));
        }

        internal void EnqueueClose(WebSocketCloseStatus status, string description)
        {
            receiveFrames.Enqueue(new ReceiveFrame(
                [], WebSocketMessageType.Close, true, status, description));
        }

        public override void Abort() => state = WebSocketState.Aborted;

        public override Task CloseAsync(
            WebSocketCloseStatus closeStatus,
            string? statusDescription,
            CancellationToken cancellationToken)
        {
            state = WebSocketState.Closed;
            return Task.CompletedTask;
        }

        public override Task CloseOutputAsync(
            WebSocketCloseStatus closeStatus,
            string? statusDescription,
            CancellationToken cancellationToken)
        {
            state = WebSocketState.CloseSent;
            return Task.CompletedTask;
        }

        public override void Dispose() => state = WebSocketState.Closed;

        public override Task<WebSocketReceiveResult> ReceiveAsync(
            ArraySegment<byte> buffer,
            CancellationToken cancellationToken)
        {
            if (!receiveFrames.TryDequeue(out var frame))
            {
                throw new InvalidOperationException("No receive frame was queued.");
            }
            Array.Copy(frame.Bytes, 0, buffer.Array!, buffer.Offset, frame.Bytes.Length);
            closeStatus = frame.CloseStatus;
            closeStatusDescription = frame.CloseDescription;
            return Task.FromResult(new WebSocketReceiveResult(
                frame.Bytes.Length,
                frame.MessageType,
                frame.EndOfMessage,
                frame.CloseStatus,
                frame.CloseDescription));
        }

        public override async Task SendAsync(
            ArraySegment<byte> buffer,
            WebSocketMessageType messageType,
            bool endOfMessage,
            CancellationToken cancellationToken)
        {
            var current = Interlocked.Increment(ref activeSends);
            InterlockedExtensions.Max(ref maximumConcurrentSends, current);
            try
            {
                if (sendDelay > TimeSpan.Zero)
                {
                    await Task.Delay(sendDelay, cancellationToken);
                }
                SentMessages.Enqueue(buffer.ToArray());
            }
            finally
            {
                Interlocked.Decrement(ref activeSends);
            }
        }

        private sealed record ReceiveFrame(
            byte[] Bytes,
            WebSocketMessageType MessageType,
            bool EndOfMessage,
            WebSocketCloseStatus? CloseStatus,
            string? CloseDescription);
    }

    private static class InterlockedExtensions
    {
        internal static void Max(ref int target, int value)
        {
            var current = Volatile.Read(ref target);
            while (current < value)
            {
                var previous = Interlocked.CompareExchange(ref target, value, current);
                if (previous == current) return;
                current = previous;
            }
        }
    }
}