From fe053cec0eb2592edaa4730b5aa8ca6e9acaedc0 Mon Sep 17 00:00:00 2001
From: Alex Batishchev <alexbat@microsoft.com>
Date: Tue, 17 May 2016 14:52:41 -0700
Subject: [PATCH] Adding exension methods ReturnsAsync and ThrowsAsync for
 ISetupSequentialResult

---
 Source/SequenceExtensions.cs           | 29 ++++++++++++++++++----
 UnitTests/SequenceExtensionsFixture.cs | 34 ++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/Source/SequenceExtensions.cs b/Source/SequenceExtensions.cs
index 98f9ed76e..7fad011ee 100644
--- a/Source/SequenceExtensions.cs
+++ b/Source/SequenceExtensions.cs
@@ -1,11 +1,8 @@
 using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Text;
+using System.Diagnostics.CodeAnalysis;
 using System.Linq.Expressions;
-using Moq.Language.Flow;
+using System.Threading.Tasks;
 using Moq.Language;
-using System.Diagnostics.CodeAnalysis;
 
 namespace Moq
 {
@@ -25,5 +22,27 @@ public static ISetupSequentialResult<TResult> SetupSequence<TMock, TResult>(
 		{
 			return new SetupSequentialContext<TMock, TResult>(mock, expression);
 		}
+
+		/// <summary>
+		/// Return a sequence of tasks, once per call.
+		/// </summary>
+		public static ISetupSequentialResult<Task<TResult>> ReturnsAsync<TResult>(this ISetupSequentialResult<Task<TResult>> setup, TResult value)
+		{
+			var tcs = new TaskCompletionSource<TResult>();
+			tcs.SetResult(value);
+
+			return setup.Returns(tcs.Task);
+		}
+
+		/// <summary>
+		/// Throws a sequence of exceptions, once per call.
+		/// </summary>
+		public static ISetupSequentialResult<Task<TResult>> ThrowsAsync<TResult>(this ISetupSequentialResult<Task<TResult>> setup, Exception exception)
+		{
+			var tcs = new TaskCompletionSource<TResult>();
+			tcs.SetException(exception);
+
+			return setup.Returns(tcs.Task);
+		}
 	}
 }
\ No newline at end of file
diff --git a/UnitTests/SequenceExtensionsFixture.cs b/UnitTests/SequenceExtensionsFixture.cs
index 091203af5..0e3a816df 100644
--- a/UnitTests/SequenceExtensionsFixture.cs
+++ b/UnitTests/SequenceExtensionsFixture.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Threading.Tasks;
 using Xunit;
 
 namespace Moq.Tests
@@ -20,6 +21,29 @@ public void PerformSequence()
 			Assert.Throws<InvalidOperationException>(() => mock.Object.Do());
 		}
 
+		[Fact]
+		public void PerformSequenceAsync()
+		{
+			var mock = new Mock<IFoo>();
+
+			mock.SetupSequence(x => x.DoAsync())
+				.ReturnsAsync(2)
+				.ReturnsAsync(3)
+				.ThrowsAsync(new InvalidOperationException());
+
+			Assert.Equal(2, mock.Object.DoAsync().Result);
+			Assert.Equal(3, mock.Object.DoAsync().Result);
+
+			try
+			{
+				var x = mock.Object.DoAsync().Result;
+			}
+			catch (AggregateException ex)
+			{
+				Assert.IsType<InvalidOperationException>(ex.GetBaseException());
+			}
+		}
+
 		[Fact]
 		public void PerformSequenceOnProperty()
 		{
@@ -67,7 +91,10 @@ public void PerformSequenceWithCallBase()
 		public interface IFoo
 		{
 			string Value { get; set; }
+
 			int Do();
+
+			Task<int> DoAsync();
 		}
 
 		public class Foo
@@ -76,6 +103,13 @@ public virtual string Do()
 			{
 				return "Ok";
 			}
+
+			public virtual Task<string> DoAsync()
+			{
+				var tcs = new TaskCompletionSource<string>();
+				tcs.SetResult("Ok");
+				return tcs.Task;
+			}
 		}
 	}
 }
\ No newline at end of file