Async overloads for OnMessage and When

src/NServiceBus.Testing.Tests/ApprovalFiles/APIApprovals.Approve.approved.txt

@@ -58,6 +58,10 @@ namespace NServiceBus.Testing
         public void OnMessage<TMessage>(System.Action<TMessage> initializeMessage = null) { }
         public void OnMessage<TMessage>(TMessage message, string messageId) { }
         public void OnMessage<TMessage>(TMessage initializedMessage) { }
+        public System.Threading.Tasks.Task OnMessageAsync<TMessage>(string messageId, System.Action<TMessage> initializeMessage = null) { }
+        public System.Threading.Tasks.Task OnMessageAsync<TMessage>(System.Action<TMessage> initializeMessage = null) { }
+        public System.Threading.Tasks.Task OnMessageAsync<TMessage>(TMessage message, string messageId) { }
+        public System.Threading.Tasks.Task OnMessageAsync<TMessage>(TMessage initializedMessage) { }
         public NServiceBus.Testing.Handler<T> SetIncomingHeader(string key, string value) { }
         public NServiceBus.Testing.Handler<T> WithExternalDependencies(System.Action<T> actionToSetUpExternalDependencies) { }
     }
@@ -139,14 +143,24 @@ namespace NServiceBus.Testing
         public NServiceBus.Testing.Saga<T> When(System.Func<T, NServiceBus.IMessageHandlerContext, System.Threading.Tasks.Task> sagaIsInvoked) { }
         public NServiceBus.Testing.Saga<T> When<TMessage>(System.Func<T, System.Func<TMessage, NServiceBus.IMessageHandlerContext, System.Threading.Tasks.Task>> handlerSelector, TMessage message) { }
         public NServiceBus.Testing.Saga<T> When<TMessage>(System.Func<T, System.Func<TMessage, NServiceBus.IMessageHandlerContext, System.Threading.Tasks.Task>> handlerSelector, System.Action<TMessage> messageInitializer = null) { }
+        public System.Threading.Tasks.Task<NServiceBus.Testing.Saga<T>> WhenAsync(System.Func<T, NServiceBus.IMessageHandlerContext, System.Threading.Tasks.Task> sagaIsInvoked) { }
+        public System.Threading.Tasks.Task<NServiceBus.Testing.Saga<T>> WhenAsync<TMessage>(System.Func<T, System.Func<TMessage, NServiceBus.IMessageHandlerContext, System.Threading.Tasks.Task>> handlerSelector, TMessage message) { }
+        public System.Threading.Tasks.Task<NServiceBus.Testing.Saga<T>> WhenAsync<TMessage>(System.Func<T, System.Func<TMessage, NServiceBus.IMessageHandlerContext, System.Threading.Tasks.Task>> handlerSelector, System.Action<TMessage> messageInitializer = null) { }
         public NServiceBus.Testing.Saga<T> WhenHandling<TMessage>(System.Action<TMessage> initializeMessage = null) { }
         public NServiceBus.Testing.Saga<T> WhenHandling<TMessage>(TMessage message) { }
+        public System.Threading.Tasks.Task<NServiceBus.Testing.Saga<T>> WhenHandlingAsync<TMessage>(System.Action<TMessage> initializeMessage = null) { }
+        public System.Threading.Tasks.Task<NServiceBus.Testing.Saga<T>> WhenHandlingAsync<TMessage>(TMessage message) { }
         public NServiceBus.Testing.Saga<T> WhenHandlingTimeout<TMessage>(System.Action<TMessage> initializeMessage = null) { }
         public NServiceBus.Testing.Saga<T> WhenHandlingTimeout<TMessage>(TMessage message) { }
+        public System.Threading.Tasks.Task<NServiceBus.Testing.Saga<T>> WhenHandlingTimeoutAsync<TMessage>(System.Action<TMessage> initializeMessage = null) { }
+        public System.Threading.Tasks.Task<NServiceBus.Testing.Saga<T>> WhenHandlingTimeoutAsync<TMessage>(TMessage message) { }
         public NServiceBus.Testing.Saga<T> WhenReceivesMessageFrom(string client) { }
         public NServiceBus.Testing.Saga<T> WhenSagaTimesOut(System.TimeSpan after) { }
         public NServiceBus.Testing.Saga<T> WhenSagaTimesOut(System.DateTime at) { }
         public NServiceBus.Testing.Saga<T> WhenSagaTimesOut() { }
+        public System.Threading.Tasks.Task<NServiceBus.Testing.Saga<T>> WhenSagaTimesOutAsync(System.TimeSpan after) { }
+        public System.Threading.Tasks.Task<NServiceBus.Testing.Saga<T>> WhenSagaTimesOutAsync(System.DateTime at) { }
+        public System.Threading.Tasks.Task<NServiceBus.Testing.Saga<T>> WhenSagaTimesOutAsync() { }
         public NServiceBus.Testing.Saga<T> WithExternalDependencies(System.Action<T> actionToSetUpExternalDependencies) { }
     }
     public class SentMessage<TMessage> : NServiceBus.Testing.OutgoingMessage<TMessage, NServiceBus.SendOptions>

src/NServiceBus.Testing/Handler.cs

@@ -1,6 +1,7 @@
 namespace NServiceBus.Testing
 {
     using System;
+    using System.Threading.Tasks;
     using MessageInterfaces.MessageMapper.Reflection;
 
     /// <summary>
@@ -266,35 +267,74 @@ public void OnMessage<TMessage>(string messageId, Action<TMessage> initializeMes
         /// </summary>
         public void OnMessage<TMessage>(Action<TMessage> initializeMessage = null)
         {
-            var message = messageCreator.CreateInstance<TMessage>();
-            initializeMessage?.Invoke(message);
-            OnMessage(message);
+            OnMessageAsync(initializeMessage).GetAwaiter().GetResult();
         }
 
         /// <summary>
         /// Activates the test that has been set up passing in given message,
         /// setting the incoming headers and the message Id.
         /// </summary>
         public void OnMessage<TMessage>(TMessage message, string messageId)
+        {
+            OnMessageAsync(message, messageId).GetAwaiter().GetResult();
+        }
+
+        /// <summary>Activates the test that has been set up passing in a specific message to be used.</summary>
+        /// <param name="initializedMessage">A message to be used with message handler.</param>
+        /// <remarks>
+        /// This is different from <see cref="OnMessage{TMessage}(System.Action{TMessage})" /> in a way that
+        /// it uses the message, and not calls to an action.
+        /// </remarks>
+        /// <example><![CDATA[var message = new TestMessage {//...}; Test.Handler<EmptyHandler>().OnMessage<TestMessage>(message);]]></example>
+        public void OnMessage<TMessage>(TMessage initializedMessage)
+        {
+            OnMessageAsync(initializedMessage).GetAwaiter().GetResult();
+        }
+
+        /// <summary>
+        /// Activates the test that has been set up passing in the given message,
+        /// setting the incoming headers and the message Id.
+        /// </summary>
+        public Task OnMessageAsync<TMessage>(string messageId, Action<TMessage> initializeMessage = null)
         {
             testableMessageHandlerContext.MessageId = messageId;
+            return OnMessageAsync(initializeMessage);
+        }
 
-            OnMessage(message);
+        /// <summary>
+        /// Activates the test that has been set up passing in the given message.
+        /// </summary>
+        public Task OnMessageAsync<TMessage>(Action<TMessage> initializeMessage = null)
+        {
+            var message = messageCreator.CreateInstance<TMessage>();
+            initializeMessage?.Invoke(message);
+            return OnMessageAsync(message);
         }
 
+        /// <summary>
+        /// Activates the test that has been set up passing in given message,
+        /// setting the incoming headers and the message Id.
+        /// </summary>
+        public Task OnMessageAsync<TMessage>(TMessage message, string messageId)
+        {
+            testableMessageHandlerContext.MessageId = messageId;
+
+            return OnMessageAsync(message);
+        }
+
         /// <summary>Activates the test that has been set up passing in a specific message to be used.</summary>
         /// <param name="initializedMessage">A message to be used with message handler.</param>
         /// <remarks>
         /// This is different from <see cref="OnMessage{TMessage}(System.Action{TMessage})" /> in a way that
         /// it uses the message, and not calls to an action.
         /// </remarks>
         /// <example><![CDATA[var message = new TestMessage {//...}; Test.Handler<EmptyHandler>().OnMessage<TestMessage>(message);]]></example>
-        public void OnMessage<TMessage>(TMessage initializedMessage)
+        public async Task OnMessageAsync<TMessage>(TMessage initializedMessage)
         {
             var messageType = messageCreator.GetMappedTypeFor(initializedMessage.GetType());
             var handleMethods = handler.GetType().CreateInvokers(messageType, typeof(IHandleMessages<>));
 
-            handleMethods.InvokeSerially(handler, initializedMessage, testableMessageHandlerContext).GetAwaiter().GetResult();
+            await handleMethods.InvokeSerially(handler, initializedMessage, testableMessageHandlerContext).ConfigureAwait(false);
 
             testableMessageHandlerContext.Validate();
         }