From 83dec1093ae3111ddca8313375dd73f4ed3d6f59 Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Fri, 13 Apr 2018 00:42:31 -0700
Subject: [PATCH] Fixed flaky longpolling tests (#1993)

- Wait for disposal and removal of the connection
---
 .../HttpConnectionContext.cs                           |  3 +++
 .../HttpConnectionDispatcher.cs                        |  5 +++--
 .../HttpConnectionDispatcherTests.cs                   | 10 +++++++---
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/Microsoft.AspNetCore.Http.Connections/HttpConnectionContext.cs b/src/Microsoft.AspNetCore.Http.Connections/HttpConnectionContext.cs
index 92a80ba310..c9ea9cc850 100644
--- a/src/Microsoft.AspNetCore.Http.Connections/HttpConnectionContext.cs
+++ b/src/Microsoft.AspNetCore.Http.Connections/HttpConnectionContext.cs
@@ -79,6 +79,9 @@ namespace Microsoft.AspNetCore.Http.Connections
 
         public SemaphoreSlim Lock { get; } = new SemaphoreSlim(1, 1);
 
+        // Used for testing only
+        internal Task DisposeAndRemoveTask { get; set; }
+
         public Task TransportTask { get; set; }
 
         public Task ApplicationTask { get; set; }
diff --git a/src/Microsoft.AspNetCore.Http.Connections/HttpConnectionDispatcher.cs b/src/Microsoft.AspNetCore.Http.Connections/HttpConnectionDispatcher.cs
index 1b0d1d501d..3327fd925b 100644
--- a/src/Microsoft.AspNetCore.Http.Connections/HttpConnectionDispatcher.cs
+++ b/src/Microsoft.AspNetCore.Http.Connections/HttpConnectionDispatcher.cs
@@ -516,8 +516,9 @@ namespace Microsoft.AspNetCore.Http.Connections
             // Complete the receiving end of the pipe
             connection.Application.Output.Complete();
 
-            // Dispose the connection gracefully, but don't wait for it
-            _ = _manager.DisposeAndRemoveAsync(connection, closeGracefully: true);
+            // Dispose the connection gracefully, but don't wait for it. We assign it here so we can wait in tests
+            connection.DisposeAndRemoveTask = _manager.DisposeAndRemoveAsync(connection, closeGracefully: true);
+
             context.Response.StatusCode = StatusCodes.Status202Accepted;
             context.Response.ContentType = "text/plain";
         }
diff --git a/test/Microsoft.AspNetCore.Http.Connections.Tests/HttpConnectionDispatcherTests.cs b/test/Microsoft.AspNetCore.Http.Connections.Tests/HttpConnectionDispatcherTests.cs
index 907ab5d225..b9d743df97 100644
--- a/test/Microsoft.AspNetCore.Http.Connections.Tests/HttpConnectionDispatcherTests.cs
+++ b/test/Microsoft.AspNetCore.Http.Connections.Tests/HttpConnectionDispatcherTests.cs
@@ -1709,13 +1709,17 @@ namespace Microsoft.AspNetCore.Http.Connections.Tests
 
                 await dispatcher.ExecuteAsync(deleteContext, options, app).OrTimeout();
 
+                // Verify the response from the DELETE request
+                Assert.Equal(StatusCodes.Status202Accepted, deleteContext.Response.StatusCode);
+                Assert.Equal("text/plain", deleteContext.Response.ContentType);
+
                 // Verify that everything shuts down
                 await connection.ApplicationTask.OrTimeout();
                 await connection.TransportTask.OrTimeout();
 
-                // Verify the response from the DELETE request
-                Assert.Equal(StatusCodes.Status202Accepted, deleteContext.Response.StatusCode);
-                Assert.Equal("text/plain", deleteContext.Response.ContentType);
+                Assert.NotNull(connection.DisposeAndRemoveTask);
+
+                await connection.DisposeAndRemoveTask.OrTimeout();
 
                 // Verify the connection was removed from the manager
                 Assert.False(manager.TryGetConnection(connection.ConnectionId, out _));