Fix flaky test (#622)

- Remove timeout and closed the application when the transport closes.
Made some tweaks to TestWebSocketConnectionFeature as a result.
- This uncovered some interesting issues with the WebSocketTransport itself so
further refactoring is needed to make it a bit more solid.

test/Microsoft.AspNetCore.Sockets.Tests/TestWebSocketConnectionFeature.cs

@@ -62,6 +62,11 @@ namespace Microsoft.AspNetCore.Sockets.Tests
                 _state = WebSocketState.Aborted;
             }
 
+            public void SendAbort()
+            {
+                _output.TryComplete(new WebSocketException(WebSocketError.ConnectionClosedPrematurely));
+            }
+
             public override async Task CloseAsync(WebSocketCloseStatus closeStatus, string statusDescription, CancellationToken cancellationToken)
             {
                 await SendMessageAsync(new WebSocketMessage
@@ -100,20 +105,39 @@ namespace Microsoft.AspNetCore.Sockets.Tests
 
             public override async Task<WebSocketReceiveResult> ReceiveAsync(ArraySegment<byte> buffer, CancellationToken cancellationToken)
             {
-                var message = await _input.ReadAsync();
+                try
-
-                if (message.MessageType == WebSocketMessageType.Close)
                 {
-                    _state = WebSocketState.CloseReceived;
+                    await _input.WaitToReadAsync();
-                    _closeStatus = message.CloseStatus;
+
-                    _closeStatusDescription = message.CloseStatusDescription;
+                    if (_input.TryRead(out var message))
-                    return new WebSocketReceiveResult(0, WebSocketMessageType.Close, true, message.CloseStatus, message.CloseStatusDescription);
+                    {
+                        if (message.MessageType == WebSocketMessageType.Close)
+                        {
+                            _state = WebSocketState.CloseReceived;
+                            _closeStatus = message.CloseStatus;
+                            _closeStatusDescription = message.CloseStatusDescription;
+                            return new WebSocketReceiveResult(0, WebSocketMessageType.Close, true, message.CloseStatus, message.CloseStatusDescription);
+                        }
+
+                        // REVIEW: This assumes the buffer passed in is > the buffer received
+                        Buffer.BlockCopy(message.Buffer, 0, buffer.Array, buffer.Offset, message.Buffer.Length);
+
+                        return new WebSocketReceiveResult(message.Buffer.Length, message.MessageType, message.EndOfMessage);
+                    }
+                }
+                catch (WebSocketException ex)
+                {
+                    switch (ex.WebSocketErrorCode)
+                    {
+                        case WebSocketError.ConnectionClosedPrematurely:
+                            _state = WebSocketState.Aborted;
+                            break;
+                    }
+
+                    throw;
                 }
 
-                // REVIEW: This assumes the buffer passed in is > the buffer received
+                throw new InvalidOperationException("Unexpected close");
-                Buffer.BlockCopy(message.Buffer, 0, buffer.Array, buffer.Offset, message.Buffer.Length);
-
-                return new WebSocketReceiveResult(message.Buffer.Length, message.MessageType, message.EndOfMessage);
             }
 
             public override Task SendAsync(ArraySegment<byte> buffer, WebSocketMessageType messageType, bool endOfMessage, CancellationToken cancellationToken)

test/Microsoft.AspNetCore.Sockets.Tests/WebSocketsTests.cs

@@ -106,12 +106,16 @@ namespace Microsoft.AspNetCore.Sockets.Tests
             using (var applicationSide = ChannelConnection.Create<byte[]>(transportToApplication, applicationToTransport))
             using (var feature = new TestWebSocketConnectionFeature())
             {
-                var options = new WebSocketOptions()
+                async Task CompleteApplicationAfterTransportCompletes()
                 {
-                    CloseTimeout = TimeSpan.FromMilliseconds(100)
+                    // Wait until the transport completes so that we can end the application
-                };
+                    await applicationSide.In.WaitToReadAsync();
 
-                var ws = new WebSocketsTransport(options, transportSide, connectionId: string.Empty, loggerFactory: new LoggerFactory());
+                    // Complete the application so that the connection unwinds without aborting
+                    applicationSide.Out.TryComplete();
+                }
+
+                var ws = new WebSocketsTransport(new WebSocketOptions(), transportSide, connectionId: string.Empty, loggerFactory: new LoggerFactory());
 
                 // Give the server socket to the transport and run it
                 var transport = ws.ProcessSocketAsync(await feature.AcceptAsync());
@@ -119,11 +123,15 @@ namespace Microsoft.AspNetCore.Sockets.Tests
                 // Run the client socket
                 var client = feature.Client.ExecuteAndCaptureFramesAsync();
 
+                // When the close frame is received, we complete the application so the send
+                // loop unwinds
+                _ = CompleteApplicationAfterTransportCompletes();
+
                 // Terminate the client to server channel with an exception
-                feature.Client.Abort();
+                feature.Client.SendAbort();
 
                 // Wait for the transport
-                await Assert.ThrowsAsync<OperationCanceledException>(() => transport).OrTimeout();
+                await Assert.ThrowsAsync<WebSocketException>(() => transport).OrTimeout();
 
                 var summary = await client.OrTimeout();
                 Assert.Equal(WebSocketCloseStatus.InternalServerError, summary.CloseResult.CloseStatus);