From 8c3b83c047c4c2dea9aa9cb7e3c8bb537bc0889e Mon Sep 17 00:00:00 2001
From: "Chris Ross (ASP.NET)" <chrross@microsoft.com>
Date: Fri, 20 Apr 2018 10:30:39 -0700
Subject: [PATCH] Fix ClientCancellationAborts test #1379

---
 .../ClientHandlerTests.cs                                 | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/test/Microsoft.AspNetCore.TestHost.Tests/ClientHandlerTests.cs b/test/Microsoft.AspNetCore.TestHost.Tests/ClientHandlerTests.cs
index 73f1c86d29..aeed94a6a8 100644
--- a/test/Microsoft.AspNetCore.TestHost.Tests/ClientHandlerTests.cs
+++ b/test/Microsoft.AspNetCore.TestHost.Tests/ClientHandlerTests.cs
@@ -12,6 +12,7 @@ using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Hosting.Server;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
+using Microsoft.AspNetCore.Testing;
 using Microsoft.AspNetCore.Testing.xunit;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
@@ -210,8 +211,8 @@ namespace Microsoft.AspNetCore.TestHost
             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);
+            var result = await readTask.TimeoutAfter(TimeSpan.FromSeconds(10));
+            Assert.Equal(0, result);
             block.Set();
         }
 
@@ -235,8 +236,7 @@ namespace Microsoft.AspNetCore.TestHost
             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());
+            await Assert.ThrowsAsync<OperationCanceledException>(() => readTask.TimeoutAfter(TimeSpan.FromSeconds(10)));
             block.Set();
         }