From 9852f6a143867c81160cfa0bd900f7562ae2bd0e Mon Sep 17 00:00:00 2001
From: Dennis Doomen <dennis.doomen@avivasolutions.nl>
Date: Fri, 7 Jun 2019 21:50:22 +0200
Subject: [PATCH] Rework on async function assertions (#1079)

---
 .../GenericAsyncFunctionAssertions.cs         |  8 ++--
 .../NonGenericAsyncFunctionAssertions.cs      | 37 ++++++++++---------
 Tests/Shared.Specs/TaskOfTAssertionSpecs.cs   |  4 +-
 3 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/Src/FluentAssertions/Specialized/GenericAsyncFunctionAssertions.cs b/Src/FluentAssertions/Specialized/GenericAsyncFunctionAssertions.cs
index 51a54f9ad2..e961e9661d 100644
--- a/Src/FluentAssertions/Specialized/GenericAsyncFunctionAssertions.cs
+++ b/Src/FluentAssertions/Specialized/GenericAsyncFunctionAssertions.cs
@@ -33,7 +33,7 @@ public GenericAsyncFunctionAssertions(Func<Task<TResult>> subject, IExtractExcep
         /// <param name="becauseArgs">
         /// Zero or more objects to format using the placeholders in <see cref="because" />.
         /// </param>
-        public AndWhichConstraint<GenericAsyncFunctionAssertions<TResult>, Task<TResult>> CompleteWithin(
+        public AndWhichConstraint<GenericAsyncFunctionAssertions<TResult>, TResult> CompleteWithin(
             TimeSpan timeSpan, string because = "", params object[] becauseArgs)
         {
             Task<TResult> task = subject();
@@ -44,7 +44,7 @@ public GenericAsyncFunctionAssertions(Func<Task<TResult>> subject, IExtractExcep
                 .BecauseOf(because, becauseArgs)
                 .FailWith("Expected {context:task} to complete within {0}{reason}.", timeSpan);
 
-            return new AndWhichConstraint<GenericAsyncFunctionAssertions<TResult>, Task<TResult>>(this, task);
+            return new AndWhichConstraint<GenericAsyncFunctionAssertions<TResult>, TResult>(this, task.Result);
         }
 
         /// <summary>
@@ -58,7 +58,7 @@ public GenericAsyncFunctionAssertions(Func<Task<TResult>> subject, IExtractExcep
         /// <param name="becauseArgs">
         /// Zero or more objects to format using the placeholders in <see cref="because" />.
         /// </param>
-        public async Task<AndWhichConstraint<GenericAsyncFunctionAssertions<TResult>, Task<TResult>>> CompleteWithinAsync(
+        public async Task<AndWhichConstraint<GenericAsyncFunctionAssertions<TResult>, TResult>> CompleteWithinAsync(
             TimeSpan timeSpan, string because = "", params object[] becauseArgs)
         {
             using (var timeoutCancellationTokenSource = new CancellationTokenSource())
@@ -79,7 +79,7 @@ public GenericAsyncFunctionAssertions(Func<Task<TResult>> subject, IExtractExcep
                     .BecauseOf(because, becauseArgs)
                     .FailWith("Expected {context:task} to complete within {0}{reason}.", timeSpan);
 
-                return new AndWhichConstraint<GenericAsyncFunctionAssertions<TResult>, Task<TResult>>(this, task);
+                return new AndWhichConstraint<GenericAsyncFunctionAssertions<TResult>, TResult>(this, task.Result);
             }
         }
     }
diff --git a/Src/FluentAssertions/Specialized/NonGenericAsyncFunctionAssertions.cs b/Src/FluentAssertions/Specialized/NonGenericAsyncFunctionAssertions.cs
index 49bbaf3abd..4fc750458c 100644
--- a/Src/FluentAssertions/Specialized/NonGenericAsyncFunctionAssertions.cs
+++ b/Src/FluentAssertions/Specialized/NonGenericAsyncFunctionAssertions.cs
@@ -32,7 +32,7 @@ public class NonGenericAsyncFunctionAssertions : AsyncFunctionAssertions
         /// <param name="becauseArgs">
         /// Zero or more objects to format using the placeholders in <see cref="because" />.
         /// </param>
-        public AndWhichConstraint<AsyncFunctionAssertions, Task> CompleteWithin(
+        public AndConstraint<AsyncFunctionAssertions> CompleteWithin(
             TimeSpan timeSpan, string because = "", params object[] becauseArgs)
         {
             Task task = Subject();
@@ -43,7 +43,7 @@ public class NonGenericAsyncFunctionAssertions : AsyncFunctionAssertions
                 .BecauseOf(because, becauseArgs)
                 .FailWith("Expected {context:task} to complete within {0}{reason}.", timeSpan);
 
-            return new AndWhichConstraint<AsyncFunctionAssertions, Task>(this, task);
+            return new AndConstraint<AsyncFunctionAssertions>(this);
         }
 
         /// <summary>
@@ -57,28 +57,29 @@ public class NonGenericAsyncFunctionAssertions : AsyncFunctionAssertions
         /// <param name="becauseArgs">
         /// Zero or more objects to format using the placeholders in <see cref="because" />.
         /// </param>
-        public async Task<AndWhichConstraint<AsyncFunctionAssertions, Task>> CompleteWithinAsync(
+        public async Task<AndConstraint<AsyncFunctionAssertions>> CompleteWithinAsync(
             TimeSpan timeSpan, string because = "", params object[] becauseArgs)
         {
-            var timeoutCancellationTokenSource = new CancellationTokenSource();
+            using (var timeoutCancellationTokenSource = new CancellationTokenSource())
+            {
+                Task task = Subject();
 
-            Task task = Subject();
+                Task completedTask =
+                    await Task.WhenAny(task, clock.DelayAsync(timeSpan, timeoutCancellationTokenSource.Token));
 
-            Task completedTask =
-                await Task.WhenAny(task, clock.DelayAsync(timeSpan, timeoutCancellationTokenSource.Token));
+                if (completedTask == task)
+                {
+                    timeoutCancellationTokenSource.Cancel();
+                    await completedTask;
+                }
 
-            if (completedTask == task)
-            {
-                timeoutCancellationTokenSource.Cancel();
-                await completedTask;
-            }
+                Execute.Assertion
+                    .ForCondition(completedTask == task)
+                    .BecauseOf(because, becauseArgs)
+                    .FailWith("Expected {context:task} to complete within {0}{reason}.", timeSpan);
 
-            Execute.Assertion
-                .ForCondition(completedTask == task)
-                .BecauseOf(because, becauseArgs)
-                .FailWith("Expected {context:task} to complete within {0}{reason}.", timeSpan);
-
-            return new AndWhichConstraint<AsyncFunctionAssertions, Task>(this, task);
+                return new AndConstraint<AsyncFunctionAssertions>(this);
+            }
         }
     }
 }
diff --git a/Tests/Shared.Specs/TaskOfTAssertionSpecs.cs b/Tests/Shared.Specs/TaskOfTAssertionSpecs.cs
index 544b8f9776..b9b78291c9 100644
--- a/Tests/Shared.Specs/TaskOfTAssertionSpecs.cs
+++ b/Tests/Shared.Specs/TaskOfTAssertionSpecs.cs
@@ -25,7 +25,7 @@ public void When_task_completes_fast_it_should_succeed()
                 Func<Task<int>> func = () => taskFactory.Task;
 
                 func.Should(timer).CompleteWithin(100.Milliseconds())
-                    .Which.Result.Should().Be(42);
+                    .Which.Should().Be(42);
             };
 
             taskFactory.SetResult(42);
@@ -81,7 +81,7 @@ public void When_task_completes_fast_async_it_should_succeed()
                 Func<Task<int>> func = () => taskFactory.Task;
 
                 (await func.Should(timer).CompleteWithinAsync(100.Milliseconds()))
-                    .Which.Result.Should().Be(42);
+                    .Which.Should().Be(42);
             };
 
             taskFactory.SetResult(42);