MoShiro/backend/tests/MoShiro.Backend.Tests/SseClient.cs

using System.Collections.Concurrent;
using System.Text.Json.Nodes;

namespace MoShiro.Backend.Tests;

public sealed record SseEvent(string Name, JsonNode? Data);

/// <summary>读取 text/event-stream：把事件和心跳注释收集起来，测试按条件等待。</summary>
public sealed class SseClient : IAsyncDisposable
{
    private readonly HttpResponseMessage _response;
    private readonly CancellationTokenSource _stop = new();
    private readonly Task _reader;

    private SseClient(HttpResponseMessage response, Stream stream)
    {
        _response = response;
        _reader = Task.Run(() => ReadAsync(stream));
    }

    public ConcurrentQueue<SseEvent> Events { get; } = new();

    private int _comments;

    public int Comments => Volatile.Read(ref _comments);

    public string? ContentType => _response.Content.Headers.ContentType?.MediaType;

    public static async Task<SseClient> ConnectAsync(HttpClient http, string url)
    {
        HttpResponseMessage response = await http.GetAsync(url, HttpCompletionOption.ResponseHeadersRead);
        response.EnsureSuccessStatusCode();
        return new SseClient(response, await response.Content.ReadAsStreamAsync());
    }

    public async Task<SseEvent> WaitForAsync(Func<SseEvent, bool> predicate, TimeSpan? timeout = null)
    {
        long started = TimeProvider.System.GetTimestamp();
        TimeSpan limit = timeout ?? TimeSpan.FromSeconds(15);
        while (TimeProvider.System.GetElapsedTime(started) < limit)
        {
            SseEvent? found = Events.ToArray().FirstOrDefault(predicate);
            if (found is not null)
            {
                return found;
            }

            await Task.Delay(25);
        }

        throw new TimeoutException("the expected server-sent event never arrived; got: " + string.Join(", ", Events.Select(e => e.Name)));
    }

    public async ValueTask DisposeAsync()
    {
        await _stop.CancelAsync();
        _response.Dispose();
        try
        {
            await _reader;
        }
        catch (Exception ex) when (ex is OperationCanceledException or IOException or ObjectDisposedException)
        {
            // 连接被关闭
        }

        _stop.Dispose();
    }

    private async Task ReadAsync(Stream stream)
    {
        using var reader = new StreamReader(stream);
        string? name = null;
        var data = new System.Text.StringBuilder();
        while (!_stop.IsCancellationRequested)
        {
            string? line = await reader.ReadLineAsync(_stop.Token);
            if (line is null)
            {
                return;
            }

            if (line.StartsWith(':'))
            {
                Interlocked.Increment(ref _comments);
            }
            else if (line.StartsWith("event: ", StringComparison.Ordinal))
            {
                name = line["event: ".Length..];
            }
            else if (line.StartsWith("data: ", StringComparison.Ordinal))
            {
                data.Append(line["data: ".Length..]);
            }
            else if (line.Length == 0 && name is not null)
            {
                Events.Enqueue(new SseEvent(name, JsonNode.Parse(data.ToString())));
                name = null;
                data.Clear();
            }
        }
    }
}