Await anything & custom awaitable types

CHANGELOG.md

@@ -7,6 +7,28 @@ The format is loosely based on [Keep a Changelog](http://keepachangelog.com/en/1
 
 ## Unreleased
 
+#### Added
+
+* "Await anything". Async methods can now be set up using the `Await` methods that act as a substitute for the `await` keyword (which the compilers don't allow inside setup expressions). This also enables method chaining across async calls in setup expressions. Should work almost anywhere.
+
+   ```csharp
+   using static Moq.AwaitOperator;
+
+   mock.SetupGet(m => Await(m.GetByIdAsync(...)).Items)
+       .Returns(new[] { 1, 2 });
+
+   var x = await mock.Object.GetByIdAsync(...);
+   var items = x.Items;  // [ 1, 2 ]
+   ```
+
+   or, using `Mock.Of`:
+
+   ```csharp
+   Mock.Of(m => Await(m.GetByIdAsync(...)).Items == new[] { 1, 2 }));
+   ```
+
+   See [details in the pull request description](https://github.com/moq/moq4/pull/1123), including an example of how to enable custom awaitable types. (@stakx, #1123)
+
 #### Changed
 
 * Attempts to mark conditionals setup as verifiable are once again allowed; it turns out that forbidding it (as was done in #997 for version 4.14.0) is in fact a regression. (@stakx, #1121)

src/Moq/ActionObserver.cs

@@ -9,6 +9,7 @@
 using System.Linq.Expressions;
 using System.Reflection;
 
+using Moq.Async;
 using Moq.Expressions.Visitors;
 using Moq.Internals;
 using Moq.Properties;
@@ -60,6 +61,19 @@ public override Expression<Action<T>> ReconstructExpression<T>(Action<T> action,
 					var invocation = recorder.Invocation;
 					if (invocation != null)
 					{
+						var resultType = invocation.Method.DeclaringType;
+						if (resultType.IsAssignableFrom(body.Type) == false)
+						{
+							if (AwaitableHandler.TryGet(body.Type) is { } awaitableHandler
+								&& awaitableHandler.ResultType.IsAssignableFrom(resultType))
+							{
+								// We are here because the current invocation cannot be chained onto the previous one,
+								// however it *can* be chained if we assume that there was an `Await` call around the
+								// former invocation that we simply don't see because static methods aren't recorded.
+								// In this case, we make things work by wrapping an `await` around the left invocation:
+								body = new AwaitExpression(body, awaitableHandler);
+							}
+						}
 						body = Expression.Call(body, invocation.Method, GetArgumentExpressions(invocation, recorder.Matches.ToArray()));
 					}
 					else
@@ -227,7 +241,7 @@ private sealed class Recorder : IInterceptor
 			private int creationTimestamp;
 			private Invocation invocation;
 			private int invocationTimestamp;
-			private IProxy returnValue;
+			private object returnValue;
 
 			public Recorder(MatcherObserver matcherObserver)
 			{
@@ -248,7 +262,7 @@ public IEnumerable<Match> Matches
 				}
 			}
 
-			public Recorder Next => this.returnValue?.Interceptor as Recorder;
+			public Recorder Next => (Unwrap.ResultIfCompletedAwaitable(this.returnValue) as IProxy)?.Interceptor as Recorder;
 
 			public void Intercept(Invocation invocation)
 			{
@@ -277,6 +291,11 @@ public void Intercept(Invocation invocation)
 					{
 						this.returnValue = null;
 					}
+					else if (AwaitableHandler.TryGet(returnType) is { } awaitableHandler)
+					{
+						var result = CreateProxy(awaitableHandler.ResultType, null, this.matcherObserver, out _);
+						this.returnValue = awaitableHandler.CreateCompleted(result);
+					}
 					else if (returnType.IsMockable())
 					{
 						this.returnValue = CreateProxy(returnType, null, this.matcherObserver, out _);

src/Moq/Async/AwaitExpression.cs

@@ -0,0 +1,41 @@
+// Copyright (c) 2007, Clarius Consulting, Manas Technology Solutions, InSTEDD, and Contributors.
+// All rights reserved. Licensed under the BSD 3-Clause License; see License.txt.
+
+using System;
+using System.Linq.Expressions;
+
+namespace Moq.Async
+{
+	/// <summary>
+	///   Represents an <see langword="await"/> expression.
+	/// </summary>
+	/// <remarks>
+	///   When reconstructing an expression tree from a delegate, <see cref="ActionObserver"/> will insert
+	///   nodes of this kind as a substitute for inferred calls to <c>Await</c> methods. (Being static methods,
+	///   they are not recorded / seen by <see cref="ActionObserver"/> and therefore their identity gets "lost".)
+	/// </remarks>
+	internal sealed class AwaitExpression : Expression
+	{
+		public readonly AwaitableHandler AwaitableHandler;
+		public readonly Expression Operand;
+
+		public AwaitExpression(Expression operand, AwaitableHandler awaitableHandler)
+		{
+			this.AwaitableHandler = awaitableHandler;
+			this.Operand = operand;
+		}
+
+		public override bool CanReduce => false;
+
+		public override ExpressionType NodeType => ExpressionType.Extension;
+
+		public override Type Type => this.AwaitableHandler.ResultType;
+
+		public override string ToString()
+		{
+			return $"(await {this.Operand})";
+		}
+
+		protected override Expression VisitChildren(ExpressionVisitor visitor) => this;
+	}
+}

src/Moq/Async/AwaitableHandler.cs

@@ -0,0 +1,123 @@
+// Copyright (c) 2007, Clarius Consulting, Manas Technology Solutions, InSTEDD, and Contributors.
+// All rights reserved. Licensed under the BSD 3-Clause License; see License.txt.
+
+using System;
+using System.Collections.Generic;
+using System.ComponentModel;
+using System.Threading.Tasks;
+
+namespace Moq.Async
+{
+	/// <summary>
+	///   Converts return values and exceptions to and from instances of a particular awaitable type.
+	/// </summary>
+	[EditorBrowsable(EditorBrowsableState.Advanced)]
+	public abstract class AwaitableHandler
+	{
+		private static readonly Dictionary<Type, Func<Type, AwaitableHandler>> factories;
+
+		static AwaitableHandler()
+		{
+			AwaitableHandler.factories = new Dictionary<Type, Func<Type, AwaitableHandler>>()
+			{
+				[typeof(Task)] = type => TaskHandler.Instance,
+				[typeof(Task<>)] = type => new TaskOfHandler(type.GetGenericArguments()[0]),
+				[typeof(ValueTask)] = type => ValueTaskHandler.Instance,
+				[typeof(ValueTask<>)] = type => new ValueTaskOfHandler(type, type.GetGenericArguments()[0]),
+			};
+		}
+
+		/// <summary>
+		/// Registers an <see cref="AwaitableHandler"/> factory function for the given awaitable type.
+		/// This allows Moq to properly recognize the awaitable type.
+		/// </summary>
+		/// <remarks>
+		/// As an example, given an <see langword="await"/>-able type <c>TaskLike&lt;TResult&gt;</c>
+		/// and a corresponding <see cref="AwaitableHandler"/> implementation <c>TaskLikeHandler</c>,
+		/// call this method as follows:
+		/// <code>
+		///   AwaitableHandler.Register(typeof(TaskLike&lt;&gt;), type => new TaskLikeHandler(resultType: type.GetGenericArguments().Single()));
+		/// </code>
+		/// </remarks>
+		/// <param name="typeDefinition">
+		///   The awaitable type for which to register the given <paramref name="factory"/> function.
+		///   You can specify an open generic type definition, e. g. <c>typeof(TaskLike&lt;&gt;)</c>,
+		///   to cover all concrete type instantiations.
+		/// </param>
+		/// <param name="factory">
+		///   The factory function that, given the <see cref="Type"/> of a concrete awaitable type,
+		///   will produce a suitable <see cref="AwaitableHandler"/> for it.</param>
+		public static void Register(Type typeDefinition, Func<Type, AwaitableHandler> factory)
+		{
+			AwaitableHandler.factories[typeDefinition] = factory;
+		}
+
+		internal static AwaitableHandler TryGet(Type type)
+		{
+			var typeDefinition = type.IsConstructedGenericType ? type.GetGenericTypeDefinition() : type;
+			return AwaitableHandler.factories.TryGetValue(typeDefinition, out var factory) ? factory.Invoke(type) : null;
+		}
+
+		/// <summary>
+		///   Gets the type of result value represented by instances of this handler's awaitable type.
+		/// </summary>
+		/// <remarks>
+		///   If this awaitable type does not have any result values, this property should return
+		///   <see langword="typeof"/>(<see langword="void"/>).
+		/// </remarks>
+		public abstract Type ResultType { get; }
+
+		/// <summary>
+		///   Converts the given result value to a successfully completed awaitable.
+		/// </summary>
+		/// <remarks>
+		///   If this awaitable types does not have any result values, <paramref name="result"/> may be ignored.
+		/// </remarks>
+		public abstract object CreateCompleted(object result);
+
+		/// <summary>
+		///   Converts the given exception to a faulted awaitable.
+		/// </summary>
+		public abstract object CreateFaulted(Exception exception);
+
+		/// <summary>
+		///   Attempts to extract the result value from the given awaitable.
+		///   This will succeed only for a successfully completed awaitable that has a result value.
+		/// </summary>
+		/// <param name="awaitable">The awaitable from which a result value should be extracted.</param>
+		/// <param name="result">
+		///   If successful, this <see langword="out"/> parameter is set to the extracted result value.
+		/// </param>
+		/// <returns>
+		///   <see langword="true"/> if extraction of a result value succeeded;
+		///   otherwise, <see langword="false"/>.
+		/// </returns>
+		public virtual bool TryGetResult(object awaitable, out object result)
+		{
+			result = null;
+
+			if (awaitable != null)
+			{
+				var awaitableType = awaitable.GetType();
+				var awaiter = awaitableType.GetMethod("GetAwaiter")?.Invoke(awaitable, null);
+				if (awaiter != null)
+				{
+					var awaiterType = awaiter.GetType();
+					var isCompleted = awaiterType.GetProperty("IsCompleted")?.GetValue(awaiter);
+					if (object.Equals(isCompleted, true))
+					{
+						try
+						{
+							result = awaiterType.GetMethod("GetResult")?.Invoke(awaiter, null);
+						}
+						catch
+						{
+						}
+					}
+				}
+			}
+
+			return result != null;
+		}
+	}
+}

src/Moq/Async/TaskHandler.cs

@@ -0,0 +1,41 @@
+// Copyright (c) 2007, Clarius Consulting, Manas Technology Solutions, InSTEDD, and Contributors.
+// All rights reserved. Licensed under the BSD 3-Clause License; see License.txt.
+
+using System;
+using System.Threading.Tasks;
+
+namespace Moq.Async
+{
+	internal sealed class TaskHandler : AwaitableHandler
+	{
+		public static readonly TaskHandler Instance = new TaskHandler();
+
+		private TaskHandler()
+		{
+		}
+
+		public override Type ResultType => typeof(void);
+
+		public object CreateCompleted()
+		{
+			var tcs = new TaskCompletionSource<bool>();
+			tcs.SetResult(true);
+			return tcs.Task;
+		}
+
+		public override object CreateCompleted(object _) => this.CreateCompleted();
+
+		public override object CreateFaulted(Exception exception)
+		{
+			var tcs = new TaskCompletionSource<bool>();
+			tcs.SetException(exception);
+			return tcs.Task;
+		}
+
+		public override bool TryGetResult(object task, out object result)
+		{
+			result = null;
+			return false;
+		}
+	}
+}

src/Moq/Async/TaskOfHandler.cs

@@ -0,0 +1,61 @@
+// Copyright (c) 2007, Clarius Consulting, Manas Technology Solutions, InSTEDD, and Contributors.
+// All rights reserved. Licensed under the BSD 3-Clause License; see License.txt.
+
+using System;
+using System.Threading.Tasks;
+
+namespace Moq.Async
+{
+	internal sealed class TaskOfHandler : AwaitableHandler
+	{
+		private readonly Type resultType;
+		private readonly Type tcsType;
+
+		public TaskOfHandler(Type resultType)
+		{
+			this.resultType = resultType;
+			this.tcsType = typeof(TaskCompletionSource<>).MakeGenericType(this.resultType);
+		}
+
+		public override Type ResultType => this.resultType;
+
+		public override object CreateCompleted(object result)
+		{
+			var tcs = Activator.CreateInstance(this.tcsType);
+			this.tcsType.GetMethod("SetResult").Invoke(tcs, new object[] { result });
+			var task = this.tcsType.GetProperty("Task").GetValue(tcs);
+			return task;
+		}
+
+		public override object CreateFaulted(Exception exception)
+		{
+			var tcs = Activator.CreateInstance(this.tcsType);
+			this.tcsType.GetMethod("SetException", new Type[] { typeof(Exception) }).Invoke(tcs, new object[] { exception });
+			var task = this.tcsType.GetProperty("Task").GetValue(tcs);
+			return task;
+		}
+
+		public override bool TryGetResult(object task, out object result)
+		{
+			if (task != null)
+			{
+				var type = task.GetType();
+				var isCompleted = (bool)type.GetProperty("IsCompleted").GetValue(task);
+				if (isCompleted)
+				{
+					try
+					{
+						result = type.GetProperty("Result").GetValue(task);
+						return true;
+					}
+					catch
+					{
+					}
+				}
+			}
+
+			result = null;
+			return false;
+		}
+	}
+}

src/Moq/Async/ValueTaskHandler.cs

@@ -0,0 +1,39 @@
+// Copyright (c) 2007, Clarius Consulting, Manas Technology Solutions, InSTEDD, and Contributors.
+// All rights reserved. Licensed under the BSD 3-Clause License; see License.txt.
+
+using System;
+using System.Threading.Tasks;
+
+namespace Moq.Async
+{
+	internal sealed class ValueTaskHandler : AwaitableHandler
+	{
+		public static readonly ValueTaskHandler Instance = new ValueTaskHandler();
+
+		private ValueTaskHandler()
+		{
+		}
+
+		public override Type ResultType => typeof(void);
+
+		public object CreateCompleted()
+		{
+			return new ValueTask();
+		}
+
+		public override object CreateCompleted(object _) => this.CreateCompleted();
+
+		public override object CreateFaulted(Exception exception)
+		{
+			var tcs = new TaskCompletionSource<bool>();
+			tcs.SetException(exception);
+			return new ValueTask(tcs.Task);
+		}
+
+		public override bool TryGetResult(object task, out object result)
+		{
+			result = null;
+			return false;
+		}
+	}
+}