Apply suggestions from code review

Co-authored-by: Miha Zupan <[email protected]>

src/libraries/System.Net.WebSockets.Client/tests/ClientWebSocketTestBase.cs

@@ -151,7 +151,7 @@ protected Task CreateEchoServerAsync(Func<Uri, Task> clientFunc)
             {
                 if (UseVersion.Major == 2)
                 {
-                    return null;
+                    throw new NotSupportedException("Remote servers don't support WebSockets over HTTP/2 yet");
                 }
 
                 if (UseSsl)

src/libraries/System.Net.WebSockets.Client/tests/SendReceiveTest.cs

@@ -642,7 +642,7 @@ public async Task SendReceive_Concurrent_Success_Base()
 
             await CreateEchoServerAsync(async uri =>
             {
-                using (ClientWebSocket cws = await GetConnectedWebSocket(uri, TimeOutMilliseconds, _output))
+                using ClientWebSocket cws = await GetConnectedWebSocket(uri, TimeOutMilliseconds, _output);
                 {
                     CancellationTokenSource ctsDefault = new CancellationTokenSource(TimeOutMilliseconds);
 
@@ -655,7 +655,7 @@ await CreateEchoServerAsync(async uri =>
 
                     for (int i = 0; i < sendBuffer.Length; i++)
                     {
-                        Task<WebSocketReceiveResult> receive = ReceiveAsync(cws, new ArraySegment<byte>(receiveBuffer, receiveBuffer.Length - i - 1, 1), ctsDefault.Token);
+                        Task<WebSocketReceiveResult> receive = ReceiveAsync(cws, new ArraySegment<byte>(receiveBuffer, i, 1), ctsDefault.Token);
                         Task send = SendAsync(cws, new ArraySegment<byte>(sendBuffer, i, 1), WebSocketMessageType.Binary, true, ctsDefault.Token);
                         await Task.WhenAll(receive, send);
                         Assert.Equal(1, receive.Result.Count);

src/libraries/System.Net.WebSockets.Client/tests/WebSocketHelper.cs

@@ -149,9 +149,9 @@ public static Task GetEchoHttp2LoopbackServer(Func<Uri, Task> clientFunc)
             return GetEchoHttp2LoopbackServer(clientFunc, new Http2Options());
         }
 
-        public static Task GetEchoHttp2LoopbackServer(Func<Uri, Task> clientFunc, Http2Options options)
+        public static async Task GetEchoHttp2LoopbackServer(Func<Uri, Task> clientFunc, Http2Options options)
         {
-            Http2LoopbackServer server = Http2LoopbackServer.CreateServer(options);
+            using Http2LoopbackServer server = Http2LoopbackServer.CreateServer(options);
 
             Task serverTask = Task.Run(async () =>
             {
@@ -170,18 +170,10 @@ public static Task GetEchoHttp2LoopbackServer(Func<Uri, Task> clientFunc, Http2O
                     var result = await websocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
                     if (result.MessageType == WebSocketMessageType.Close)
                     {
-                        if (result.CloseStatus == WebSocketCloseStatus.Empty)
-                        {
-                            await websocket.CloseAsync(WebSocketCloseStatus.Empty, null, CancellationToken.None);
-                        }
-                        else
-                        {
-                            WebSocketCloseStatus closeStatus = result.CloseStatus.GetValueOrDefault();
-                            await websocket.CloseAsync(
-                                closeStatus,
-                                result.CloseStatusDescription,
-                                CancellationToken.None);
-                        }
+                        await websocket.CloseAsync(
+                            result.CloseStatus ?? WebSocketCloseStatus.Empty,
+                            result.CloseStatusDescription,
+                            CancellationToken.None);
 
                         continue;
                     }
@@ -192,13 +184,13 @@ await websocket.CloseAsync(
                 await connection.DisposeAsync();
             });
 
-            return new Task[] { serverTask, clientFunc(server.Address) }.WhenAllOrAnyFailed(TimeOutMilliseconds * 2);
+            await new Task[] { serverTask, clientFunc(server.Address) }.WhenAllOrAnyFailed(TimeOutMilliseconds * 2);
         }
 
-        public static Task GetEchoLoopbackServer(Func<Uri, Task> clientFunc,  LoopbackServer.Options options = null)
+        public static async Task GetEchoLoopbackServer(Func<Uri, Task> clientFunc,  LoopbackServer.Options options = null)
         {
-            LoopbackServer server = new LoopbackServer(options);
-            Task.WaitAll(server.ListenAsync());
+            using LoopbackServer server = new LoopbackServer(options);
+            await server.ListenAsync();
 
             Task serverTask = server.AcceptConnectionAsync(async connection =>
             {
@@ -210,18 +202,10 @@ public static Task GetEchoLoopbackServer(Func<Uri, Task> clientFunc,  LoopbackSe
                     var result = await websocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
                     if (result.MessageType == WebSocketMessageType.Close)
                     {
-                        if (result.CloseStatus == WebSocketCloseStatus.Empty)
-                        {
-                            await websocket.CloseAsync(WebSocketCloseStatus.Empty, null, CancellationToken.None);
-                        }
-                        else
-                        {
-                            WebSocketCloseStatus closeStatus = result.CloseStatus.GetValueOrDefault();
-                            await websocket.CloseAsync(
-                                closeStatus,
-                                result.CloseStatusDescription,
-                                CancellationToken.None);
-                        }
+                        await websocket.CloseAsync(
+                            result.CloseStatus ?? WebSocketCloseStatus.Empty,
+                            result.CloseStatusDescription,
+                            CancellationToken.None);
 
                         continue;
                     }
@@ -233,7 +217,7 @@ await websocket.CloseAsync(
                 await connection.DisposeAsync();
             });
 
-            return new Task[] { serverTask, clientFunc(server.Address) }.WhenAllOrAnyFailed(TimeOutMilliseconds * 2); ;
+            await new Task[] { serverTask, clientFunc(server.Address) }.WhenAllOrAnyFailed(TimeOutMilliseconds * 2); ;
         }
 
         private static bool InitWebSocketSupported()