From 9363eff2a8261c511923c74065bbbfc7e12c2278 Mon Sep 17 00:00:00 2001
From: Dylan Dmitri Gray <d.dylan.g@gmail.com>
Date: Fri, 17 May 2019 10:40:19 -0700
Subject: [PATCH] OrTimeout Extension

---
 .../RequestThrottling/sample/Startup.cs         |  2 ++
 .../test/SemaphoreWrapperTests.cs               | 17 ++++++++++++-----
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/Middleware/RequestThrottling/sample/Startup.cs b/src/Middleware/RequestThrottling/sample/Startup.cs
index a9bc74e319..ea4969c9bd 100644
--- a/src/Middleware/RequestThrottling/sample/Startup.cs
+++ b/src/Middleware/RequestThrottling/sample/Startup.cs
@@ -26,7 +26,9 @@ namespace RequestThrottlingSample
         {
             app.Run(async context =>
             {
+                Console.WriteLine("HEWWWO?");
                 await context.Response.WriteAsync("Hello world!");
+                Console.WriteLine("GOODBWWYWE!");
             });
         }
 
diff --git a/src/Middleware/RequestThrottling/test/SemaphoreWrapperTests.cs b/src/Middleware/RequestThrottling/test/SemaphoreWrapperTests.cs
index 368804c8cd..f9ddd80420 100644
--- a/src/Middleware/RequestThrottling/test/SemaphoreWrapperTests.cs
+++ b/src/Middleware/RequestThrottling/test/SemaphoreWrapperTests.cs
@@ -10,6 +10,13 @@ using Microsoft.AspNetCore.Testing;
 
 namespace Microsoft.AspNetCore.RequestThrottling.Tests
 {
+    public static class TaskExtensions
+    {
+        public static Task OrTimeout(this Task task, int seconds = 30)
+        {
+            return task.TimeoutAfter(TimeSpan.FromSeconds(seconds));
+        }
+    }
     public class SemaphoreWrapperTests
     {
         [Fact]
@@ -18,7 +25,7 @@ namespace Microsoft.AspNetCore.RequestThrottling.Tests
             using var s = new SemaphoreWrapper(1);
             Assert.Equal(1, s.Count);
 
-            await s.EnterQueue();
+            await s.EnterQueue().OrTimeout();
             Assert.Equal(0, s.Count);
 
             s.LeaveQueue();
@@ -44,13 +51,13 @@ namespace Microsoft.AspNetCore.RequestThrottling.Tests
         public async Task WaitsIfNoSpaceAvailible()
         {
             using var s = new SemaphoreWrapper(1);
-            await s.EnterQueue();
+            await s.EnterQueue().OrTimeout();
 
             var waitingTask = s.EnterQueue();
             Assert.False(waitingTask.IsCompleted);
 
             s.LeaveQueue();
-            await waitingTask.TimeoutAfter(TimeSpan.FromSeconds(30));
+            await waitingTask.OrTimeout();
         }
 
         [Fact]
@@ -59,8 +66,8 @@ namespace Microsoft.AspNetCore.RequestThrottling.Tests
             using var s1 = new SemaphoreWrapper(1);
             using var s2 = new SemaphoreWrapper(1);
 
-            await s1.EnterQueue();
-            await s2.EnterQueue().TimeoutAfter(TimeSpan.FromSeconds(30));
+            await s1.EnterQueue().OrTimeout();
+            await s2.EnterQueue().OrTimeout();
         }
     }
 }