Await anything & custom awaitable types

src/Moq/ExpressionExtensions.cs

@@ -257,6 +257,13 @@ void Split(Expression e, out Expression r /* remainder */, out InvocationShape p
 						else if (methodCallExpression.IsAwait(out var awaitableHandler))
 						{
 							Split(methodCallExpression.Arguments.Single(), out r, out p);
+							p = new InvocationShape(
+								expression:
+									Expression.Lambda(
+										new AwaitExpression(methodCallExpression.Arguments.Single(), awaitableHandler),
+										p.Expression.Parameters),
+								method: p.Method,
+								arguments: p.Arguments);
 							p.AwaitableHandler = awaitableHandler;
 						}
 						else

tests/Moq.Tests/AwaitOperatorFixture.cs

@@ -2,6 +2,7 @@
 // All rights reserved. Licensed under the BSD 3-Clause License; see License.txt.
 
 using System;
+using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Threading.Tasks;
 
@@ -589,6 +590,29 @@ public async Task Mock_Of()
 			Assert.Equal(expectedSameName, actualName);
 		}
 
+		[Fact]
+		public void Await__is_represented__in_setup_Expression()
+		{
+			var mock = new Mock<IPerson>();
+			mock.Setup(m => Await(m.DoSomethingTaskAsync()));
+			var setup = mock.Setups.Last();
+			Assert.Contains("aWAiT", setup.Expression.ToStringFixed(), StringComparison.OrdinalIgnoreCase);
+			//               ^^^^^
+			// NOTE: Depending on the exact scenario, the `Await` method name may get converted to
+			// a faux `await` keyword (e.g. with `SetupSet`). Right now, we don't care too much about
+			// that slight loss of accuracy, as long as a clue for the occurred await remains.
+		}
+
+		[Fact]
+		public void Await__is_represented__in_setup_OriginalExpression()
+		{
+			// NOTE: The note in the test above applies here, too.
+			var mock = new Mock<IPerson>();
+			mock.Setup(m => Await(m.DoSomethingTaskAsync()));
+			var setup = mock.Setups.Last();
+			Assert.Contains("await", setup.OriginalExpression.ToStringFixed(), StringComparison.OrdinalIgnoreCase);
+		}
+
 		public interface IPerson
 		{
 			IPerson Friend { get; }