Deflake Http2WebSocket tests (#48742)

Author: BrennanConroy

src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2WebSocketTests.cs

@@ -155,7 +155,6 @@ await InitializeConnectionAsync(async context =>
     [QuarantinedTest("https://github.com/dotnet/aspnetcore/issues/42133")]
     public async Task HEADERS_Received_SecondRequest_Accepted()
     {
-        // Add stream to Http2Connection._completedStreams inline with SetResult().
         var appDelegateTcs = new TaskCompletionSource();
         await InitializeConnectionAsync(async context =>
         {
@@ -181,6 +180,19 @@ await InitializeConnectionAsync(async context =>
             await appDelegateTcs.Task;
         });
 
+        var originalHandler = _connection._streamLifetimeHandler;
+        var tcs = new TaskCompletionSource();
+        var streamLifetimeHandler = new Mock<IHttp2StreamLifetimeHandler>();
+        streamLifetimeHandler.Setup(o => o.OnStreamCompleted(It.IsAny<Http2Stream>())).Callback((Http2Stream stream) =>
+        {
+            // Add stream to Http2Connection._completedStreams.
+            originalHandler.OnStreamCompleted(stream);
+
+            // Unblock test code that will call TriggerTick and return the stream to the pool
+            tcs.TrySetResult();
+        });
+        _connection._streamLifetimeHandler = streamLifetimeHandler.Object;
+
         // HEADERS + END_HEADERS
         // :method = CONNECT
         // :protocol = websocket
@@ -224,6 +236,7 @@ await InitializeConnectionAsync(async context =>
         Assert.Equal(0x01, dataFrame.Payload.Span[0]);
 
         appDelegateTcs.TrySetResult();
+        await tcs.Task.DefaultTimeout();
 
         dataFrame = await ExpectAsync(Http2FrameType.DATA,
             withLength: 0,
@@ -595,8 +608,23 @@ await InitializeConnectionAsync(async context =>
                 // We've done the test. Now just return the normal echo server behavior.
                 await _echoApplication(context);
             }
+        });
 
+        var originalHandler = _connection._streamLifetimeHandler;
+        var tcs = new TaskCompletionSource();
+        var streamLifetimeHandler = new Mock<IHttp2StreamLifetimeHandler>();
+        streamLifetimeHandler.Setup(o => o.OnStreamCompleted(It.IsAny<Http2Stream>())).Callback((Http2Stream stream) =>
+        {
+            // Add stream to Http2Connection._completedStreams.
+            originalHandler.OnStreamCompleted(stream);
+
+            if (requestCount == 1)
+            {
+                // Unblock test code that will call TriggerTick and return the stream to the pool
+                tcs.TrySetResult();
+            }
         });
+        _connection._streamLifetimeHandler = streamLifetimeHandler.Object;
 
         // HEADERS + END_HEADERS
         // :method = CONNECT
@@ -636,6 +664,8 @@ await InitializeConnectionAsync(async context =>
         Assert.Equal(0x01, dataFrame.Payload.Span[0]);
 
         appDelegateTcs.TrySetResult();
+        await tcs.Task.DefaultTimeout();
+
         dataFrame = await ExpectAsync(Http2FrameType.DATA,
             withLength: 0,
             withFlags: (byte)Http2DataFrameFlags.END_STREAM,