From c8bf886e92f411d77aa2412cb42519a28c88f441 Mon Sep 17 00:00:00 2001
From: BrennanConroy <brecon@microsoft.com>
Date: Mon, 22 May 2017 15:16:35 -0700
Subject: [PATCH] Fix tests not running as expected (#1080)

---
 .../WebHostTests.cs                           | 24 ++++++++++++-------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/test/Microsoft.AspNetCore.Hosting.Tests/WebHostTests.cs b/test/Microsoft.AspNetCore.Hosting.Tests/WebHostTests.cs
index 888ba0cdcd..298f2d85a5 100644
--- a/test/Microsoft.AspNetCore.Hosting.Tests/WebHostTests.cs
+++ b/test/Microsoft.AspNetCore.Hosting.Tests/WebHostTests.cs
@@ -203,10 +203,12 @@ namespace Microsoft.AspNetCore.Hosting
 
             var server = new Mock<IServer>();
             server.Setup(s => s.StopAsync(It.IsAny<CancellationToken>()))
-                .Returns(Task.CompletedTask)
-                .Callback<CancellationToken>(token =>
+                .Returns<CancellationToken>(token =>
                 {
-                    token.WaitHandle.WaitOne();
+                    return Task.Run(() =>
+                    {
+                        token.WaitHandle.WaitOne();
+                    });
                 });
 
             using (var host = CreateBuilder(config)
@@ -240,10 +242,12 @@ namespace Microsoft.AspNetCore.Hosting
 
             var server = new Mock<IServer>();
             server.Setup(s => s.StopAsync(It.IsAny<CancellationToken>()))
-                .Returns(Task.CompletedTask)
-                .Callback<CancellationToken>(token =>
+                .Returns<CancellationToken>(token =>
                 {
-                    token.WaitHandle.WaitOne();
+                    return Task.Run(() =>
+                    {
+                        token.WaitHandle.WaitOne();
+                    });
                 });
 
             using (var host = CreateBuilder(config)
@@ -274,10 +278,12 @@ namespace Microsoft.AspNetCore.Hosting
 
             var server = new Mock<IServer>();
             server.Setup(s => s.StopAsync(It.IsAny<CancellationToken>()))
-                .Returns(Task.CompletedTask)
-                .Callback<CancellationToken>(token =>
+                .Returns<CancellationToken>(token =>
                 {
-                    token.WaitHandle.WaitOne();
+                    return Task.Run(() =>
+                    {
+                        token.WaitHandle.WaitOne();
+                    });
                 });
 
             using (var host = CreateBuilder(config)