Replace ManualResetEvent with TaskCompletionSource in TestHost tests

Author: Tratcher

src/Hosting/TestHost/test/ClientHandlerTests.cs

@@ -156,51 +156,50 @@ public async Task BlockingMiddlewareShouldNotBlockClient()
         [Fact]
         public async Task HeadersAvailableBeforeBodyFinished()
         {
-            ManualResetEvent block = new ManualResetEvent(false);
+            var block = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var handler = new ClientHandler(PathString.Empty, new DummyApplication(async context =>
             {
                 context.Response.Headers["TestHeader"] = "TestValue";
                 await context.Response.WriteAsync("BodyStarted,");
-                block.WaitOne();
+                await block.Task;
                 await context.Response.WriteAsync("BodyFinished");
             }));
             var httpClient = new HttpClient(handler);
             HttpResponseMessage response = await httpClient.GetAsync("https://example.com/",
                 HttpCompletionOption.ResponseHeadersRead);
             Assert.Equal("TestValue", response.Headers.GetValues("TestHeader").First());
-            block.Set();
+            block.SetResult(0);
             Assert.Equal("BodyStarted,BodyFinished", await response.Content.ReadAsStringAsync());
         }
 
         [Fact]
         public async Task FlushSendsHeaders()
         {
-            ManualResetEvent block = new ManualResetEvent(false);
+            var block = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var handler = new ClientHandler(PathString.Empty, new DummyApplication(async context =>
             {
                 context.Response.Headers["TestHeader"] = "TestValue";
                 context.Response.Body.Flush();
-                block.WaitOne();
+                await block.Task;
                 await context.Response.WriteAsync("BodyFinished");
             }));
             var httpClient = new HttpClient(handler);
             HttpResponseMessage response = await httpClient.GetAsync("https://example.com/",
                 HttpCompletionOption.ResponseHeadersRead);
             Assert.Equal("TestValue", response.Headers.GetValues("TestHeader").First());
-            block.Set();
+            block.SetResult(0);
             Assert.Equal("BodyFinished", await response.Content.ReadAsStringAsync());
         }
 
         [Fact]
         public async Task ClientDisposalCloses()
         {
-            ManualResetEvent block = new ManualResetEvent(false);
+            var block = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var handler = new ClientHandler(PathString.Empty, new DummyApplication(context =>
             {
                 context.Response.Headers["TestHeader"] = "TestValue";
                 context.Response.Body.Flush();
-                block.WaitOne();
-                return Task.FromResult(0);
+                return block.Task;
             }));
             var httpClient = new HttpClient(handler);
             HttpResponseMessage response = await httpClient.GetAsync("https://example.com/",
@@ -210,21 +209,20 @@ public async Task ClientDisposalCloses()
             Task<int> readTask = responseStream.ReadAsync(new byte[100], 0, 100);
             Assert.False(readTask.IsCompleted);
             responseStream.Dispose();
-            Assert.True(readTask.Wait(TimeSpan.FromSeconds(10)), "Finished");
-            Assert.Equal(0, readTask.Result);
-            block.Set();
+            var read = await readTask.WithTimeout();
+            Assert.Equal(0, read);
+            block.SetResult(0);
         }
 
         [Fact]
         public async Task ClientCancellationAborts()
         {
-            ManualResetEvent block = new ManualResetEvent(false);
+            var block = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var handler = new ClientHandler(PathString.Empty, new DummyApplication(context =>
             {
                 context.Response.Headers["TestHeader"] = "TestValue";
                 context.Response.Body.Flush();
-                block.WaitOne();
-                return Task.FromResult(0);
+                return block.Task;
             }));
             var httpClient = new HttpClient(handler);
             HttpResponseMessage response = await httpClient.GetAsync("https://example.com/",
@@ -235,9 +233,8 @@ public async Task ClientCancellationAborts()
             Task<int> readTask = responseStream.ReadAsync(new byte[100], 0, 100, cts.Token);
             Assert.False(readTask.IsCompleted, "Not Completed");
             cts.Cancel();
-            var ex = Assert.Throws<AggregateException>(() => readTask.Wait(TimeSpan.FromSeconds(10)));
-            Assert.IsAssignableFrom<OperationCanceledException>(ex.GetBaseException());
-            block.Set();
+            await Assert.ThrowsAsync<OperationCanceledException>(() => readTask.WithTimeout());
+            block.SetResult(0);
         }
 
         [Fact]
@@ -255,19 +252,19 @@ public Task ExceptionBeforeFirstWriteIsReported()
         [Fact]
         public async Task ExceptionAfterFirstWriteIsReported()
         {
-            ManualResetEvent block = new ManualResetEvent(false);
+            var block = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var handler = new ClientHandler(PathString.Empty, new DummyApplication(async context =>
             {
                 context.Response.Headers["TestHeader"] = "TestValue";
                 await context.Response.WriteAsync("BodyStarted");
-                block.WaitOne();
+                await block.Task;
                 throw new InvalidOperationException("Test Exception");
             }));
             var httpClient = new HttpClient(handler);
             HttpResponseMessage response = await httpClient.GetAsync("https://example.com/",
                 HttpCompletionOption.ResponseHeadersRead);
             Assert.Equal("TestValue", response.Headers.GetValues("TestHeader").First());
-            block.Set();
+            block.SetResult(0);
             var ex = await Assert.ThrowsAsync<HttpRequestException>(() => response.Content.ReadAsStringAsync());
             Assert.IsType<InvalidOperationException>(ex.GetBaseException());
         }

src/Hosting/TestHost/test/HttpContextBuilderTests.cs

@@ -102,18 +102,17 @@ public async Task BlockingMiddlewareShouldNotBlockClient()
         [Fact]
         public async Task HeadersAvailableBeforeSyncBodyFinished()
         {
-            var block = new ManualResetEvent(false);
+            var block = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var builder = new WebHostBuilder().Configure(app =>
             {
-                app.Run(c =>
+                app.Run(async c =>
                 {
                     c.Response.Headers["TestHeader"] = "TestValue";
                     var bytes = Encoding.UTF8.GetBytes("BodyStarted" + Environment.NewLine);
                     c.Response.Body.Write(bytes, 0, bytes.Length);
-                    Assert.True(block.WaitOne(TimeSpan.FromSeconds(5)));
+                    await block.Task;
                     bytes = Encoding.UTF8.GetBytes("BodyFinished");
                     c.Response.Body.Write(bytes, 0, bytes.Length);
-                    return Task.CompletedTask;
                 });
             });
             var server = new TestServer(builder);
@@ -122,21 +121,21 @@ public async Task HeadersAvailableBeforeSyncBodyFinished()
             Assert.Equal("TestValue", context.Response.Headers["TestHeader"]);
             var reader = new StreamReader(context.Response.Body);
             Assert.Equal("BodyStarted", reader.ReadLine());
-            block.Set();
+            block.SetResult(0);
             Assert.Equal("BodyFinished", reader.ReadToEnd());
         }
 
         [Fact]
         public async Task HeadersAvailableBeforeAsyncBodyFinished()
         {
-            var block = new ManualResetEvent(false);
+            var block = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var builder = new WebHostBuilder().Configure(app =>
             {
                 app.Run(async c =>
                 {
                     c.Response.Headers["TestHeader"] = "TestValue";
                     await c.Response.WriteAsync("BodyStarted" + Environment.NewLine);
-                    Assert.True(block.WaitOne(TimeSpan.FromSeconds(5)));
+                    await block.Task;
                     await c.Response.WriteAsync("BodyFinished");
                 });
             });
@@ -146,43 +145,43 @@ public async Task HeadersAvailableBeforeAsyncBodyFinished()
             Assert.Equal("TestValue", context.Response.Headers["TestHeader"]);
             var reader = new StreamReader(context.Response.Body);
             Assert.Equal("BodyStarted", await reader.ReadLineAsync());
-            block.Set();
+            block.SetResult(0);
             Assert.Equal("BodyFinished", await reader.ReadToEndAsync());
         }
 
         [Fact]
         public async Task FlushSendsHeaders()
         {
-            var block = new ManualResetEvent(false);
+            var block = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var builder = new WebHostBuilder().Configure(app =>
             {
                 app.Run(async c =>
                 {
                     c.Response.Headers["TestHeader"] = "TestValue";
                     c.Response.Body.Flush();
-                    block.WaitOne();
+                    await block.Task;
                     await c.Response.WriteAsync("BodyFinished");
                 });
             });
             var server = new TestServer(builder);
             var context = await server.SendAsync(c => { });
 
             Assert.Equal("TestValue", context.Response.Headers["TestHeader"]);
-            block.Set();
+            block.SetResult(0);
             Assert.Equal("BodyFinished", new StreamReader(context.Response.Body).ReadToEnd());
         }
 
         [Fact]
         public async Task ClientDisposalCloses()
         {
-            var block = new ManualResetEvent(false);
+            var block = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var builder = new WebHostBuilder().Configure(app =>
             {
                 app.Run(async c =>
                 {
                     c.Response.Headers["TestHeader"] = "TestValue";
                     c.Response.Body.Flush();
-                    block.WaitOne();
+                    await block.Task;
                     await c.Response.WriteAsync("BodyFinished");
                 });
             });
@@ -194,20 +193,20 @@ public async Task ClientDisposalCloses()
             Task<int> readTask = responseStream.ReadAsync(new byte[100], 0, 100);
             Assert.False(readTask.IsCompleted);
             responseStream.Dispose();
-            Assert.True(readTask.Wait(TimeSpan.FromSeconds(10)));
-            Assert.Equal(0, readTask.Result);
-            block.Set();
+            var read = await readTask.WithTimeout();
+            Assert.Equal(0, read);
+            block.SetResult(0);
         }
 
         [Fact]
-        public void ClientCancellationAborts()
+        public async Task ClientCancellationAborts()
         {
-            var block = new ManualResetEvent(false);
+            var block = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var builder = new WebHostBuilder().Configure(app =>
             {
                 app.Run(c =>
                 {
-                    block.Set();
+                    block.SetResult(0);
                     Assert.True(c.RequestAborted.WaitHandle.WaitOne(TimeSpan.FromSeconds(10)));
                     c.RequestAborted.ThrowIfCancellationRequested();
                     return Task.CompletedTask;
@@ -216,24 +215,23 @@ public void ClientCancellationAborts()
             var server = new TestServer(builder);
             var cts = new CancellationTokenSource();
             var contextTask = server.SendAsync(c => { }, cts.Token);
-            block.WaitOne();
+            await block.Task;
             cts.Cancel();
 
-            var ex = Assert.Throws<AggregateException>(() => contextTask.Wait(TimeSpan.FromSeconds(10)));
-            Assert.IsAssignableFrom<OperationCanceledException>(ex.GetBaseException());
+            await Assert.ThrowsAsync<OperationCanceledException>(() => contextTask.WithTimeout());
         }
 
         [Fact]
         public async Task ClientCancellationAbortsReadAsync()
         {
-            var block = new ManualResetEvent(false);
+            var block = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var builder = new WebHostBuilder().Configure(app =>
             {
                 app.Run(async c =>
                 {
                     c.Response.Headers["TestHeader"] = "TestValue";
                     c.Response.Body.Flush();
-                    block.WaitOne();
+                    await block.Task;
                     await c.Response.WriteAsync("BodyFinished");
                 });
             });
@@ -246,9 +244,8 @@ public async Task ClientCancellationAbortsReadAsync()
             var readTask = responseStream.ReadAsync(new byte[100], 0, 100, cts.Token);
             Assert.False(readTask.IsCompleted);
             cts.Cancel();
-            var ex = Assert.Throws<AggregateException>(() => readTask.Wait(TimeSpan.FromSeconds(10)));
-            Assert.IsAssignableFrom<OperationCanceledException>(ex.GetBaseException());
-            block.Set();
+            await Assert.ThrowsAsync<OperationCanceledException>(() => readTask.WithTimeout());
+            block.SetResult(0);
         }
 
         [Fact]
@@ -268,14 +265,14 @@ public Task ExceptionBeforeFirstWriteIsReported()
         [Fact]
         public async Task ExceptionAfterFirstWriteIsReported()
         {
-            var block = new ManualResetEvent(false);
+            var block = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var builder = new WebHostBuilder().Configure(app =>
             {
                 app.Run(async c =>
                 {
                     c.Response.Headers["TestHeader"] = "TestValue";
                     await c.Response.WriteAsync("BodyStarted");
-                    block.WaitOne();
+                    await block.Task;
                     throw new InvalidOperationException("Test Exception");
                 });
             });
@@ -284,7 +281,7 @@ public async Task ExceptionAfterFirstWriteIsReported()
 
             Assert.Equal("TestValue", context.Response.Headers["TestHeader"]);
             Assert.Equal(11, context.Response.Body.Read(new byte[100], 0, 100));
-            block.Set();
+            block.SetResult(0);
             var ex = Assert.Throws<IOException>(() => context.Response.Body.Read(new byte[100], 0, 100));
             Assert.IsAssignableFrom<InvalidOperationException>(ex.InnerException);
         }

src/Hosting/TestHost/test/Utilities.cs

@@ -0,0 +1,29 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Threading.Tasks;
+
+namespace Microsoft.AspNetCore.TestHost
+{
+    internal static class Utilities
+    {
+        internal static readonly TimeSpan DefaultTimeout = TimeSpan.FromSeconds(15);
+
+        internal static Task<T> WithTimeout<T>(this Task<T> task) => task.WithTimeout(DefaultTimeout);
+
+        internal static async Task<T> WithTimeout<T>(this Task<T> task, TimeSpan timeout)
+        {
+            var completedTask = await Task.WhenAny(task, Task.Delay(timeout));
+
+            if (completedTask == task)
+            {
+                return await task;
+            }
+            else
+            {
+                throw new TimeoutException("The task has timed out.");
+            }
+        }
+    }
+}