From 0655213fd67d9265a00300b245c33650963bb282 Mon Sep 17 00:00:00 2001
From: Artur Krajewski <krajek.dev@gmail.com>
Date: Thu, 4 Oct 2018 21:00:18 +0200
Subject: [PATCH] Add `ExecutionTime` acception async action

---
 Src/FluentAssertions/AssertionExtensions.cs   | 17 +++++++-
 .../Specialized/ExecutionTimeAssertions.cs    | 37 ++++++++++++++++
 .../ExecutionTimeAssertionsSpecs.cs           | 43 ++++++++++++++++++-
 3 files changed, 94 insertions(+), 3 deletions(-)

diff --git a/Src/FluentAssertions/AssertionExtensions.cs b/Src/FluentAssertions/AssertionExtensions.cs
index 6ff58aa00f..2523875328 100644
--- a/Src/FluentAssertions/AssertionExtensions.cs
+++ b/Src/FluentAssertions/AssertionExtensions.cs
@@ -55,9 +55,9 @@ public static MemberExecutionTime<T> ExecutionTimeOf<T>(this T subject, Expressi
         }
 
         /// <summary>
-        /// Provides methods for asserting the execution time of a method or property.
+        /// Provides methods for asserting the execution time of an action.
         /// </summary>
-        /// <param name="action">A reference to the method or property to measure the execution time of.</param>
+        /// <param name="action">An action to measure the execution time of.</param>
         /// <returns>
         /// Returns an object for asserting that the execution time matches certain conditions.
         /// </returns>
@@ -67,6 +67,19 @@ public static ExecutionTime ExecutionTime(this Action action)
             return new ExecutionTime(action);
         }
 
+        /// <summary>
+        /// Provides methods for asserting the execution time of an async action.
+        /// </summary>
+        /// <param name="action">An async action to measure the execution time of.</param>
+        /// <returns>
+        /// Returns an object for asserting that the execution time matches certain conditions.
+        /// </returns>
+        [MustUseReturnValue /* do not use Pure because this method executes the action before returning to the caller */]
+        public static ExecutionTime ExecutionTime(this Func<Task> action)
+        {
+            return new ExecutionTime(action);
+        }
+
         /// <summary>
         /// Returns an <see cref="ExecutionTimeAssertions"/> object that can be used to assert the
         /// current <see cref="ExecutionTime"/>.
diff --git a/Src/FluentAssertions/Specialized/ExecutionTimeAssertions.cs b/Src/FluentAssertions/Specialized/ExecutionTimeAssertions.cs
index 55de9ea488..2fcbe81be7 100644
--- a/Src/FluentAssertions/Specialized/ExecutionTimeAssertions.cs
+++ b/Src/FluentAssertions/Specialized/ExecutionTimeAssertions.cs
@@ -220,6 +220,11 @@ public ExecutionTime(Action action)
         {
         }
 
+        public ExecutionTime(Func<Task> action)
+            : this(action, "the action")
+        {
+        }
+
         protected ExecutionTime(Action action, string actionDescription)
         {
             ActionDescription = actionDescription;
@@ -247,6 +252,38 @@ protected ExecutionTime(Action action, string actionDescription)
             });
         }
 
+        /// <remarks>
+        /// This constructor is almost exact copy of the one accepting <see cref="Action"/>.
+        /// The original constructor shall stay in place in order to keep backward-compatibility
+        /// and to avoid unnecessary wrapping action in <see cref="Task"/>.
+        /// </remarks>
+        protected ExecutionTime(Func<Task> action, string actionDescription)
+        {
+            ActionDescription = actionDescription;
+            stopwatch = new Stopwatch();
+            IsRunning = true;
+            Task = Task.Run(async () =>
+            {
+                // move stopwatch as close to action start as possible
+                // so that we have to get correct time readings
+                try
+                {
+                    stopwatch.Start();
+                    await action();
+                }
+                catch (Exception exception)
+                {
+                    Exception = exception;
+                }
+                finally
+                {
+                    // ensures that we stop the stopwatch even on exceptions
+                    stopwatch.Stop();
+                    IsRunning = false;
+                }
+            });
+        }
+
         internal TimeSpan ElapsedTime => stopwatch.Elapsed;
 
         internal bool IsRunning { get; private set; }
diff --git a/Tests/Shared.Specs/ExecutionTimeAssertionsSpecs.cs b/Tests/Shared.Specs/ExecutionTimeAssertionsSpecs.cs
index dffb3e7f76..c1666b12b6 100644
--- a/Tests/Shared.Specs/ExecutionTimeAssertionsSpecs.cs
+++ b/Tests/Shared.Specs/ExecutionTimeAssertionsSpecs.cs
@@ -1,6 +1,6 @@
 using System;
 using System.Threading;
-
+using System.Threading.Tasks;
 using FluentAssertions.Extensions;
 using Xunit;
 using Xunit.Sdk;
@@ -179,6 +179,27 @@ public void When_the_execution_time_of_an_action_is_not_less_than_a_limit__it_sh
                 "*action should be less than 0.100s, but it required*");
         }
 
+        [Fact]
+        public void When_the_execution_time_of_an_async_action_is_not_less_than_a_limit__it_should_throw()
+        {
+            //-----------------------------------------------------------------------------------------------------------
+            // Arrange
+            //-----------------------------------------------------------------------------------------------------------
+
+            Func<Task> someAction = async () => await Task.Delay(TimeSpan.FromMilliseconds(150));
+
+            //-----------------------------------------------------------------------------------------------------------
+            // Act
+            //-----------------------------------------------------------------------------------------------------------
+            Action act = () => someAction.ExecutionTime().Should().BeLessThan(100.Milliseconds());
+
+            //-----------------------------------------------------------------------------------------------------------
+            // Assert
+            //-----------------------------------------------------------------------------------------------------------
+            act.Should().Throw<XunitException>().WithMessage(
+                "*action should be less than 0.100s, but it required*");
+        }
+
         [Fact]
         public void When_the_execution_time_of_an_action_is_less_than_a_limit_it_should_not_throw()
         {
@@ -199,6 +220,26 @@ public void When_the_execution_time_of_an_action_is_less_than_a_limit_it_should_
             act.Should().NotThrow();
         }
 
+        [Fact]
+        public void When_the_execution_time_of_an_async_action_is_less_than_a_limit_it_should_not_throw()
+        {
+            //-----------------------------------------------------------------------------------------------------------
+            // Arrange
+            //-----------------------------------------------------------------------------------------------------------
+
+            Func<Task> someAction = async () => await Task.Delay(TimeSpan.FromMilliseconds(100));
+
+            //-----------------------------------------------------------------------------------------------------------
+            // Act
+            //-----------------------------------------------------------------------------------------------------------
+            Action act = () => someAction.ExecutionTime().Should().BeLessThan(1.Seconds());
+
+            //-----------------------------------------------------------------------------------------------------------
+            // Assert
+            //-----------------------------------------------------------------------------------------------------------
+            act.Should().NotThrow();
+        }
+
         [Fact]
         public void When_action_runs_indefinitely_it_should_be_stopped_and_throw_if_there_is_less_than_condition()
         {