WithTimeout extension methods for use in UnitTests

WithTimeoutExtensions.cs

public static class WithTimeoutExtension
{
    public static Task WithTimeout(this Task task, int timeout, string actionName = null) => task.WithTimeout(TimeSpan.FromMilliseconds(timeout), actionName);

    public static async Task WithTimeout(this Task task, TimeSpan timeout, string actionName = null)
    {
        var delay = Task.Delay(timeout);
        var completedTask = await Task.WhenAny(task, delay);
        if (completedTask == task)
        {
            return;
        }

        Assert.Fail($"{actionName ?? "Task"} is not completed in {timeout} for {TestContext.CurrentContext.CurrentRepeatCount + 1} try");
    }

    public static async Task<T> WithTimeout<T>(this Task<T> task, TimeSpan timeout, string actionName = null)
    {
        var delay = Task.Delay(timeout);
        var completedTask = await Task.WhenAny(task, delay);
        if (completedTask == task)
        {
            return task.GetAwaiter().GetResult();
        }

        throw new AssertionException($"{actionName ?? "Task"} is not completed in {timeout} for {TestContext.CurrentContext.CurrentRepeatCount + 1} try");
    }
}