From 0c0687e9ed443e7c41f928562ddd4b4b0531941c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lukas=20Gr=C3=BCtzmacher?= <github@lg2.de>
Date: Wed, 17 Aug 2022 11:56:00 +0200
Subject: [PATCH] SAVEPOINT

---
 .../Specialized/AsyncFunctionAssertions.cs    | 61 +++++++++++++++++--
 .../Specialized/DelegateAssertions.cs         |  2 +-
 .../Specialized/DelegateAssertionsBase.cs     | 19 ++++--
 3 files changed, 70 insertions(+), 12 deletions(-)

diff --git a/Src/FluentAssertions/Specialized/AsyncFunctionAssertions.cs b/Src/FluentAssertions/Specialized/AsyncFunctionAssertions.cs
index 32471cecdb..529fe59bf1 100644
--- a/Src/FluentAssertions/Specialized/AsyncFunctionAssertions.cs
+++ b/Src/FluentAssertions/Specialized/AsyncFunctionAssertions.cs
@@ -132,7 +132,7 @@ public AsyncFunctionAssertions(Func<TTask> subject, IExtractExceptions extractor
             .BecauseOf(because, becauseArgs)
             .FailWith("Expected {context} to throw exactly {0}{reason}, but found <null>.", expectedType);
 
-        Exception exception = await InvokeWithInterceptionAsync(Subject);
+        Exception exception = await InvokeWithInterceptionAsync(Subject, TimeSpan.MaxValue);
 
         Execute.Assertion
             .ForCondition(exception is not null)
@@ -163,8 +163,33 @@ public AsyncFunctionAssertions(Func<TTask> subject, IExtractExceptions extractor
             .BecauseOf(because, becauseArgs)
             .FailWith("Expected {context} to throw {0}{reason}, but found <null>.", typeof(TException));
 
-        Exception exception = await InvokeWithInterceptionAsync(Subject);
-        return ThrowInternal<TException>(exception, because, becauseArgs);
+        Exception exception = await InvokeWithInterceptionAsync(Subject, TimeSpan.MaxValue);
+        return ThrowInternal<TException>(exception, TimeSpan.MaxValue, because, becauseArgs);
+    }
+
+    /// <summary>
+    /// Asserts that the current <see cref="Func{Task}"/> throws an exception of type <typeparamref name="TException"/>.
+    /// </summary>
+    /// <param name="timeSpan">The allowed time span for the operation.</param>
+    /// <param name="because">
+    /// A formatted phrase as is supported by <see cref="string.Format(string,object[])" /> explaining why the assertion
+    /// is needed. If the phrase does not start with the word <i>because</i>, it is prepended automatically.
+    /// </param>
+    /// <param name="becauseArgs">
+    /// Zero or more objects to format using the placeholders in <paramref name="because" />.
+    /// </param>
+    public async Task<ExceptionAssertions<TException>> ThrowWithinAsync<TException>(
+        TimeSpan timeSpan, string because = "", params object[] becauseArgs)
+        where TException : Exception
+    {
+        Execute.Assertion
+            .ForCondition(Subject is not null)
+            .BecauseOf(because, becauseArgs)
+            .FailWith("Expected {context} to throw {0} within {1}{reason}, but found <null>.",
+                typeof(TException), timeSpan);
+
+        Exception exception = await InvokeWithInterceptionAsync(Subject, timeSpan);
+        return ThrowInternal<TException>(exception, timeSpan, because, becauseArgs);
     }
 
     /// <summary>
@@ -277,7 +302,7 @@ async Task<AndConstraint<TAssertions>> AssertionTaskAsync()
 
             while (invocationEndTime is null || invocationEndTime < waitTime)
             {
-                exception = await InvokeWithInterceptionAsync(Subject);
+                exception = await InvokeWithInterceptionAsync(Subject, TimeSpan.MaxValue);
                 if (exception is null)
                 {
                     return new AndConstraint<TAssertions>((TAssertions)this);
@@ -327,7 +352,7 @@ private protected async Task<bool> CompletesWithinTimeoutAsync(Task target, Time
         return true;
     }
 
-    private static async Task<Exception> InvokeWithInterceptionAsync(Func<Task> action)
+    private async Task<Exception> InvokeWithInterceptionAsync(Func<Task> action, TimeSpan timeSpan)
     {
         try
         {
@@ -343,7 +368,31 @@ private static async Task<Exception> InvokeWithInterceptionAsync(Func<Task> acti
                     ? CallerIdentifier.OverrideStackSearchUsingCurrentScope()
                     : default)
             {
-                await action();
+                if (timeSpan == TimeSpan.MaxValue)
+                {
+                    await action();
+                }
+                else
+                {
+                    (TTask result, TimeSpan remainingTime) = InvokeWithTimer(timeSpan);
+                    if (remainingTime < TimeSpan.Zero)
+                    {
+                        // timeout reached without exception
+                        return null;
+                    }
+
+                    if (result.IsFaulted)
+                    {
+                        // exception in synchronous portion
+                        return result.Exception;
+                    }
+
+                    if (await CompletesWithinTimeoutAsync(result, remainingTime))
+                    {
+                        // completed without ex exception
+                        return null;
+                    }
+                }
             }
 
             return null;
diff --git a/Src/FluentAssertions/Specialized/DelegateAssertions.cs b/Src/FluentAssertions/Specialized/DelegateAssertions.cs
index 26d23a8b03..53db756e7b 100644
--- a/Src/FluentAssertions/Specialized/DelegateAssertions.cs
+++ b/Src/FluentAssertions/Specialized/DelegateAssertions.cs
@@ -45,7 +45,7 @@ public ExceptionAssertions<TException> Throw<TException>(string because = "", pa
 
         FailIfSubjectIsAsyncVoid();
         Exception exception = InvokeSubjectWithInterception();
-        return ThrowInternal<TException>(exception, because, becauseArgs);
+        return ThrowInternal<TException>(exception, TimeSpan.MaxValue, because, becauseArgs);
     }
 
     /// <summary>
diff --git a/Src/FluentAssertions/Specialized/DelegateAssertionsBase.cs b/Src/FluentAssertions/Specialized/DelegateAssertionsBase.cs
index 37b7e18cbb..87748d36bb 100644
--- a/Src/FluentAssertions/Specialized/DelegateAssertionsBase.cs
+++ b/Src/FluentAssertions/Specialized/DelegateAssertionsBase.cs
@@ -12,7 +12,9 @@ namespace FluentAssertions.Specialized;
 /// Contains a number of methods to assert that a method yields the expected result.
 /// </summary>
 [DebuggerNonUserCode]
-public abstract class DelegateAssertionsBase<TDelegate, TAssertions> : ReferenceTypeAssertions<TDelegate, DelegateAssertionsBase<TDelegate, TAssertions>>
+public abstract class
+    DelegateAssertionsBase<TDelegate, TAssertions> : ReferenceTypeAssertions<TDelegate,
+        DelegateAssertionsBase<TDelegate, TAssertions>>
     where TDelegate : Delegate
     where TAssertions : DelegateAssertionsBase<TDelegate, TAssertions>
 {
@@ -27,14 +29,21 @@ private protected DelegateAssertionsBase(TDelegate @delegate, IExtractExceptions
 
     private protected IClock Clock { get; }
 
-    protected ExceptionAssertions<TException> ThrowInternal<TException>(Exception exception, string because, object[] becauseArgs)
+    protected ExceptionAssertions<TException> ThrowInternal<TException>(
+        Exception exception, TimeSpan timeSpan, string because, object[] becauseArgs)
         where TException : Exception
     {
         TException[] expectedExceptions = extractor.OfType<TException>(exception).ToArray();
 
-        Execute.Assertion
-            .BecauseOf(because, becauseArgs)
-            .WithExpectation("Expected a <{0}> to be thrown{reason}, ", typeof(TException))
+        AssertionScope becauseOf = Execute.Assertion
+            .BecauseOf(because, becauseArgs);
+        AssertionScope expectation =
+            timeSpan == TimeSpan.MaxValue
+                ? becauseOf.WithExpectation("Expected a <{0}> to be thrown{reason}, ",
+                    typeof(TException))
+                : becauseOf.WithExpectation("Expected a <{0}> to be thrown within {1}{reason}, ",
+                    typeof(TException), timeSpan);
+        expectation
             .ForCondition(exception is not null)
             .FailWith("but no exception was thrown.")
             .Then