From 9004c960064de87fdacd24540e4c378d9655ecb6 Mon Sep 17 00:00:00 2001
From: Dennis Doomen <dennis.doomen@avivasolutions.nl>
Date: Thu, 28 Nov 2024 19:20:09 +0100
Subject: [PATCH] Fixed a regression in which CompleteWithinAsync treated a
 canceled task as an exception.

---
 .../Specialized/AsyncFunctionAssertions.cs    | 15 ++++++++------
 .../GenericAsyncFunctionAssertions.cs         |  2 +-
 .../Specialized/TaskAssertionSpecs.cs         | 20 +++++++++++++++++++
 docs/_pages/releases.md                       |  4 ++++
 4 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/Src/FluentAssertions/Specialized/AsyncFunctionAssertions.cs b/Src/FluentAssertions/Specialized/AsyncFunctionAssertions.cs
index 2a80ef257e..9fe7162408 100644
--- a/Src/FluentAssertions/Specialized/AsyncFunctionAssertions.cs
+++ b/Src/FluentAssertions/Specialized/AsyncFunctionAssertions.cs
@@ -60,7 +60,7 @@ public async Task<AndConstraint<TAssertions>> CompleteWithinAsync(
 
             if (success)
             {
-                bool completesWithinTimeout = await CompletesWithinTimeoutAsync(task, remainingTime);
+                bool completesWithinTimeout = await CompletesWithinTimeoutAsync(task, remainingTime, _ => Task.CompletedTask);
 
                 Execute.Assertion
                     .ForCondition(completesWithinTimeout)
@@ -97,7 +97,7 @@ public async Task<AndConstraint<TAssertions>> NotCompleteWithinAsync(
 
             if (remainingTime >= TimeSpan.Zero)
             {
-                bool completesWithinTimeout = await CompletesWithinTimeoutAsync(task, remainingTime);
+                bool completesWithinTimeout = await CompletesWithinTimeoutAsync(task, remainingTime, _ => Task.CompletedTask);
 
                 Execute.Assertion
                     .ForCondition(!completesWithinTimeout)
@@ -271,7 +271,11 @@ private async Task<Exception> InvokeWithInterceptionAsync(TimeSpan timeout)
                 // Here we do not need to know whether the task completes (successfully) in timeout
                 // or does not complete. We are only interested in the exception which is thrown, not returned.
                 // So, we can ignore the result.
-                _ = await CompletesWithinTimeoutAsync(task, remainingTime);
+                _ = await CompletesWithinTimeoutAsync(task, remainingTime, async cancelledTask =>
+                {
+                    // Rethrow the exception causing the task be canceled.
+                    await cancelledTask;
+                });
             }
 
             return null;
@@ -431,7 +435,7 @@ private protected (TTask result, TimeSpan remainingTime) InvokeWithTimer(TimeSpa
     /// <summary>
     ///     Monitors the specified task whether it completes withing the remaining time span.
     /// </summary>
-    private protected async Task<bool> CompletesWithinTimeoutAsync(Task target, TimeSpan remainingTime)
+    private protected async Task<bool> CompletesWithinTimeoutAsync(Task target, TimeSpan remainingTime, Func<Task, Task> onTaskCanceled)
     {
         using var delayCancellationTokenSource = new CancellationTokenSource();
 
@@ -452,8 +456,7 @@ private protected async Task<bool> CompletesWithinTimeoutAsync(Task target, Time
 
         if (target.IsCanceled)
         {
-            // Rethrow the exception causing the task be canceled.
-            await target;
+            await onTaskCanceled(target);
         }
 
         // The monitored task is completed, we shall cancel the clock.
diff --git a/Src/FluentAssertions/Specialized/GenericAsyncFunctionAssertions.cs b/Src/FluentAssertions/Specialized/GenericAsyncFunctionAssertions.cs
index d389955d0e..65dc8260cb 100644
--- a/Src/FluentAssertions/Specialized/GenericAsyncFunctionAssertions.cs
+++ b/Src/FluentAssertions/Specialized/GenericAsyncFunctionAssertions.cs
@@ -51,7 +51,7 @@ public GenericAsyncFunctionAssertions(Func<Task<TResult>> subject, IExtractExcep
 
             if (success)
             {
-                bool completesWithinTimeout = await CompletesWithinTimeoutAsync(task, remainingTime);
+                bool completesWithinTimeout = await CompletesWithinTimeoutAsync(task, remainingTime, _ => Task.CompletedTask);
 
                 success = Execute.Assertion
                     .ForCondition(completesWithinTimeout)
diff --git a/Tests/FluentAssertions.Specs/Specialized/TaskAssertionSpecs.cs b/Tests/FluentAssertions.Specs/Specialized/TaskAssertionSpecs.cs
index dcab1926ff..c71e5cba2f 100644
--- a/Tests/FluentAssertions.Specs/Specialized/TaskAssertionSpecs.cs
+++ b/Tests/FluentAssertions.Specs/Specialized/TaskAssertionSpecs.cs
@@ -104,6 +104,26 @@ public async Task When_task_completes_late_it_should_fail()
             // Assert
             await action.Should().ThrowAsync<XunitException>();
         }
+
+        [Fact]
+        public async Task Cancelling_a_token_also_means_the_task_is_completed()
+        {
+            // Arrange
+            var timer = new FakeClock();
+            var taskFactory = new TaskCompletionSource<bool>();
+
+            // Act
+            Func<Task> action = () => taskFactory
+                .Awaiting(t => (Task)t.Task)
+                .Should(timer)
+                .CompleteWithinAsync(100.Milliseconds());
+
+            taskFactory.SetCanceled();
+            timer.Complete();
+
+            // Assert
+            await action.Should().NotThrowAsync();
+        }
     }
 
     public class NotCompleteWithinAsync
diff --git a/docs/_pages/releases.md b/docs/_pages/releases.md
index 7da0b45ee8..60bcad50d8 100644
--- a/docs/_pages/releases.md
+++ b/docs/_pages/releases.md
@@ -9,6 +9,10 @@ sidebar:
 
 ## 7.0.0
 
+### Fixes
+
+* Fixed a regression in which `CompleteWithinAsync` treated a canceled task as an exception - [#9999](https://github.com/fluentassertions/fluentassertions/pull/9999)
+
 ### Breaking Changes
 
 * Dropped direct support for .NET Core 2.x and .NET Core 3.x - [#2302](https://github.com/fluentassertions/fluentassertions/pull/2302)