mxaccessgw/src/ZB.MOM.WW.MxGateway.Tests/TestSupport/RecordingServerStreamWriter.cs

using Grpc.Core;

namespace ZB.MOM.WW.MxGateway.Tests.TestSupport;

/// <summary>
/// Thread-safe <see cref="IServerStreamWriter{T}"/> that records every written message
/// and lets a test await the first message or a specific message count with a timeout.
/// </summary>
/// <typeparam name="T">The streamed message type.</typeparam>
public sealed class RecordingServerStreamWriter<T> : IServerStreamWriter<T>
{
    private readonly object _syncRoot = new();
    private readonly TaskCompletionSource<T> _firstMessage = new(TaskCreationOptions.RunContinuationsAsynchronously);
    private readonly List<T> _messages = [];
    private readonly List<TaskCompletionSource<IReadOnlyList<T>>> _countWaiters = [];

    /// <summary>Gets the messages written to this stream, in order.</summary>
    public IReadOnlyList<T> Messages
    {
        get
        {
            lock (_syncRoot)
            {
                return _messages.ToArray();
            }
        }
    }

    /// <summary>Gets or sets options for writing messages to the stream.</summary>
    public WriteOptions? WriteOptions { get; set; }

    /// <summary>Records the supplied message.</summary>
    /// <param name="message">The message to record.</param>
    /// <returns>A completed task.</returns>
    public Task WriteAsync(T message)
    {
        List<TaskCompletionSource<IReadOnlyList<T>>>? satisfied = null;
        IReadOnlyList<T>? snapshot = null;

        lock (_syncRoot)
        {
            _messages.Add(message);
            _firstMessage.TrySetResult(message);

            // Check whether any count waiters are now satisfied.
            if (_countWaiters.Count > 0)
            {
                snapshot = _messages.ToArray();
                satisfied = _countWaiters.ToList();
                _countWaiters.Clear();
            }
        }

        if (satisfied is not null && snapshot is not null)
        {
            foreach (TaskCompletionSource<IReadOnlyList<T>> waiter in satisfied)
            {
                waiter.TrySetResult(snapshot);
            }
        }

        return Task.CompletedTask;
    }

    /// <summary>Waits for the first message to be written within the specified timeout.</summary>
    /// <param name="timeout">Maximum time to wait for the first message.</param>
    /// <returns>The first message written to this stream.</returns>
    public async Task<T> WaitForFirstMessageAsync(TimeSpan timeout) =>
        await _firstMessage.Task.WaitAsync(timeout).ConfigureAwait(false);

    /// <summary>
    /// Waits until at least <paramref name="count"/> messages have been written, then returns
    /// the current snapshot. The wait is bounded by a single deadline of <c>now + timeout</c>;
    /// intermediate wakeups (when a message arrives but the count is not yet met) consume from
    /// that same deadline so the total elapsed time never exceeds <paramref name="timeout"/>.
    /// If fewer than <paramref name="count"/> messages arrive before the deadline the call
    /// throws <see cref="OperationCanceledException"/>.
    /// </summary>
    /// <param name="count">Minimum number of messages to wait for.</param>
    /// <param name="timeout">Maximum total time to wait, measured from the moment of the call.</param>
    /// <returns>A snapshot of all messages received so far (at least <paramref name="count"/>).</returns>
    public async Task<IReadOnlyList<T>> WaitForMessageCountAsync(int count, TimeSpan timeout)
    {
        // Capture a single deadline so every iteration of the loop below draws from the
        // same budget — using WaitAsync(timeout) per iteration would reset the clock on
        // each intermediate wakeup, effectively giving N×timeout total budget.
        using CancellationTokenSource deadlineCts = new(timeout);
        CancellationToken deadlineToken = deadlineCts.Token;

        TaskCompletionSource<IReadOnlyList<T>>? tcs = null;

        lock (_syncRoot)
        {
            if (_messages.Count >= count)
            {
                return _messages.ToArray();
            }

            tcs = new TaskCompletionSource<IReadOnlyList<T>>(TaskCreationOptions.RunContinuationsAsynchronously);
            _countWaiters.Add(tcs);
        }

        // Re-check each time any message arrives.  The TCS is satisfied on every write,
        // but the caller may need more messages, so we loop until the count is met.
        while (true)
        {
            IReadOnlyList<T> snapshot = await tcs.Task.WaitAsync(deadlineToken).ConfigureAwait(false);
            if (snapshot.Count >= count)
            {
                return snapshot;
            }

            // Not enough yet — register a new waiter and keep waiting against the same deadline.
            lock (_syncRoot)
            {
                if (_messages.Count >= count)
                {
                    return _messages.ToArray();
                }

                tcs = new TaskCompletionSource<IReadOnlyList<T>>(TaskCreationOptions.RunContinuationsAsynchronously);
                _countWaiters.Add(tcs);
            }
        }
    }
}