riojoy/tests/RioJoy.Core.Tests/TestPolyfills.cs

namespace RioJoy.Core.Tests;

/// <summary>
/// net48 test-only polyfills. Visible to all tests via enclosing-namespace scope.
/// </summary>
internal static class TaskTestExtensions
{
    /// <summary>
    /// Polyfill for Task&lt;T&gt;.WaitAsync(TimeSpan) (net6+), absent on net48.
    /// Throws <see cref="TimeoutException"/> if the task does not complete in time.
    /// </summary>
    public static async Task<T> WaitAsync<T>(this Task<T> task, TimeSpan timeout)
    {
        using var cts = new CancellationTokenSource();
        Task completed = await Task.WhenAny(task, Task.Delay(timeout, cts.Token)).ConfigureAwait(false);
        if (completed != task)
            throw new TimeoutException($"Task did not complete within {timeout}.");
        cts.Cancel(); // stop the delay timer
        return await task.ConfigureAwait(false);
    }

    /// <summary>Non-generic counterpart of <see cref="WaitAsync{T}"/>.</summary>
    public static async Task WaitAsync(this Task task, TimeSpan timeout)
    {
        using var cts = new CancellationTokenSource();
        Task completed = await Task.WhenAny(task, Task.Delay(timeout, cts.Token)).ConfigureAwait(false);
        if (completed != task)
            throw new TimeoutException($"Task did not complete within {timeout}.");
        cts.Cancel();
        await task.ConfigureAwait(false);
    }
}