gremlin-dotnet/test/Gremlin.Net.UnitTest/Driver/ConnectionTests.cs - tinkerpop - Git at Google

 #region License

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */

 #endregion

 using System;
 using System.Collections.Generic;
 using System.Net.WebSockets;
 using System.Threading;
 using System.Threading.Tasks;
 using Gremlin.Net.Driver;
 using Gremlin.Net.Driver.Exceptions;
 using Gremlin.Net.Driver.Messages;
 using Gremlin.Net.Structure.IO.GraphBinary;
 using Moq;
 using Xunit;

 namespace Gremlin.Net.UnitTest.Driver
 {
     public class ConnectionTests
     {
         [Fact]
         public async Task ShouldHandleCloseMessageAfterConnectAsync()
         {
             var mockedClientWebSocket = new Mock<IClientWebSocket>();
             mockedClientWebSocket
                 .Setup(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()))
                 .ReturnsAsync(new WebSocketReceiveResult(0, WebSocketMessageType.Close, true, WebSocketCloseStatus.MessageTooBig, "Message is too large"));
             mockedClientWebSocket
                 .SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);

             Uri uri = new Uri("wss://localhost:8182");
             Connection connection = GetConnection(mockedClientWebSocket, uri: uri);

             await connection.ConnectAsync(CancellationToken.None);

             Assert.False(connection.IsOpen);
             Assert.Equal(0, connection.NrRequestsInFlight);
             mockedClientWebSocket.Verify(m => m.ConnectAsync(uri, It.IsAny<CancellationToken>()), Times.Once);
             mockedClientWebSocket.Verify(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()), Times.Once);
             mockedClientWebSocket.Verify(m => m.CloseAsync(WebSocketCloseStatus.NormalClosure, string.Empty, It.IsAny<CancellationToken>()), Times.Once);
         }

         [Fact]
         public async Task ShouldThrowIfClosedMessageReceivedWithValidPropertiesAsync()
         {
             var mockedClientWebSocket = new Mock<IClientWebSocket>();
             mockedClientWebSocket
                 .SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);

             WebSocketConnection webSocketConnection = new WebSocketConnection(
                 mockedClientWebSocket.Object,
                 new WebSocketSettings());

             // Test all known close statuses
             foreach (Enum closeStatus in Enum.GetValues(typeof(WebSocketCloseStatus)))
             {
                 mockedClientWebSocket
                     .Setup(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()))
                     .ReturnsAsync(new WebSocketReceiveResult(0, WebSocketMessageType.Close, true, (WebSocketCloseStatus)closeStatus, closeStatus.ToString()));

                 await AssertExpectedConnectionClosedException((WebSocketCloseStatus?)closeStatus, closeStatus.ToString(), () => webSocketConnection.ReceiveMessageAsync());
             }

             // Test null/empty close property values as well.
             mockedClientWebSocket
                 .Setup(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()))
                 .ReturnsAsync(new WebSocketReceiveResult(0, WebSocketMessageType.Close, true, null, null));
             await AssertExpectedConnectionClosedException(null, null, () => webSocketConnection.ReceiveMessageAsync());

             mockedClientWebSocket
                 .Setup(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()))
                 .ReturnsAsync(new WebSocketReceiveResult(0, WebSocketMessageType.Close, true, null, String.Empty));
             await AssertExpectedConnectionClosedException(null, String.Empty, () => webSocketConnection.ReceiveMessageAsync());
         }

         [Fact]
         public async Task ShouldThrowOnSubmitRequestIfWebSocketIsNotOpen()
         {
             // This is to test that race bugs don't get introduced which would cause submitted requests to hang if the
             // connection is being closed/aborted or is not connected. In these cases, WebSocket.SendAsync should throw for the underlying
             // websocket is not open and the caller should be notified of that failure.
             var mockedClientWebSocket = new Mock<IClientWebSocket>();

             mockedClientWebSocket
                 .SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);

             Connection connection = GetConnection(mockedClientWebSocket);
             RequestMessage request = RequestMessage.Build("gremlin").Create();

             // Simulate the SendAsync exception behavior if the underlying websocket is closed (see reference https://docs.microsoft.com/en-us/dotnet/api/system.net.websockets.clientwebsocket.sendasync?view=net-6.0)
             mockedClientWebSocket.Setup(m => m.SendAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()))
                 .ThrowsAsync(new ObjectDisposedException(nameof(ClientWebSocket), "Socket closed"));

             // Test various closing/closed WebSocketStates with SubmitAsync.
             mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.Closed);
             await Assert.ThrowsAsync<ObjectDisposedException>(() => connection.SubmitAsync<dynamic>(request, CancellationToken.None));

             mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.CloseSent);
             await Assert.ThrowsAsync<ObjectDisposedException>(() => connection.SubmitAsync<dynamic>(request, CancellationToken.None));

             mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.CloseReceived);
             await Assert.ThrowsAsync<ObjectDisposedException>(() => connection.SubmitAsync<dynamic>(request, CancellationToken.None));

             mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.Aborted);
             await Assert.ThrowsAsync<ObjectDisposedException>(() => connection.SubmitAsync<dynamic>(request, CancellationToken.None));

             // Simulate SendAsync exception behavior if underlying websocket is not connected.
             mockedClientWebSocket.Setup(m => m.SendAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()))
                 .ThrowsAsync(new InvalidOperationException("Socket not connected"));
             mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.Connecting);
             await Assert.ThrowsAsync<InvalidOperationException>(() => connection.SubmitAsync<dynamic>(request, CancellationToken.None));
         }

         [Fact]
         public async Task ShouldHandleCloseMessageForInFlightRequestsAsync()
         {
             // Tests that in-flight requests will get notified if a connection close message is received.
             Uri uri = new Uri("wss://localhost:8182");
             WebSocketReceiveResult closeResult = new WebSocketReceiveResult(0, WebSocketMessageType.Close, true, WebSocketCloseStatus.EndpointUnavailable, "Server shutdown");

             var receiveSempahore = new SemaphoreSlim(0, 1);
             var mockedClientWebSocket = new Mock<IClientWebSocket>();
             mockedClientWebSocket
                 .Setup(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()))
                 .Returns(async () =>
                 {
                     await receiveSempahore.WaitAsync();
                     mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.CloseReceived);
                     return closeResult;
                 });
             mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.Open);
             mockedClientWebSocket
                 .SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);

             var connection = GetConnection(mockedClientWebSocket, uri: uri);
             await connection.ConnectAsync(CancellationToken.None);

             // Create two in-flight requests that will block on waiting for a response.
             RequestMessage requestMsg1 = RequestMessage.Build("gremlin").Create();
             RequestMessage requestMsg2 = RequestMessage.Build("gremlin").Create();
             Task request1 = connection.SubmitAsync<dynamic>(requestMsg1, CancellationToken.None);
             Task request2 = connection.SubmitAsync<dynamic>(requestMsg2, CancellationToken.None);

             // Confirm the requests are in-flight.
             Assert.Equal(2, connection.NrRequestsInFlight);

             // Release the connection close message.
             receiveSempahore.Release();

             // Assert that both requests get notified with the closed exception.
             await AssertExpectedConnectionClosedException(closeResult.CloseStatus, closeResult.CloseStatusDescription, () => request1);
             await AssertExpectedConnectionClosedException(closeResult.CloseStatus, closeResult.CloseStatusDescription, () => request2);

             // delay for NotifyAboutConnectionFailure running in another thread
             var runs = 0;
             while (connection.NrRequestsInFlight == 2 && ++runs < 5)
             {
                 await Task.Delay(5);
             }

             Assert.False(connection.IsOpen);
             Assert.Equal(0, connection.NrRequestsInFlight);
             mockedClientWebSocket.Verify(m => m.ConnectAsync(uri, It.IsAny<CancellationToken>()), Times.Once);
             mockedClientWebSocket.Verify(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()), Times.Once);
             mockedClientWebSocket.Verify(m => m.CloseAsync(WebSocketCloseStatus.NormalClosure, string.Empty, It.IsAny<CancellationToken>()), Times.Once);
         }

         [Fact]
         public async Task ShouldProperlyHandleCancellationForSubmitAsync()
         {
             var mockedClientWebSocket = new Mock<IClientWebSocket>();
             mockedClientWebSocket
                 .SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);
             var connection = GetConnection(mockedClientWebSocket);
             var cts = new CancellationTokenSource();

             var task = connection.SubmitAsync<object>(RequestMessage.Build(string.Empty).Create(),
                 cts.Token);
             cts.Cancel();

             await Assert.ThrowsAsync<TaskCanceledException>(async () => await task);
             Assert.True(task.IsCanceled);
             mockedClientWebSocket.Verify(m => m.SendAsync(It.IsAny<ArraySegment<byte>>(),
                 It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(), cts.Token), Times.Once);
             mockedClientWebSocket.Verify(
                 m => m.CloseAsync(It.IsAny<WebSocketCloseStatus>(), It.IsAny<string>(), It.IsAny<CancellationToken>()),
                 Times.Never);
             mockedClientWebSocket.Verify(
                 m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), cts.Token), Times.Never);
         }

         [Fact]
         public async Task ShouldProperlyHandleCancellationForSubmitAsyncIfAlreadyCancelled()
         {
             var mockedClientWebSocket = new Mock<IClientWebSocket>();
             mockedClientWebSocket
                 .SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);
             var connection = GetConnection(mockedClientWebSocket);
             var token = new CancellationToken(canceled: true);

             var task = connection.SubmitAsync<object>(RequestMessage.Build(string.Empty).Create(),
                 token);

             await Assert.ThrowsAsync<TaskCanceledException>(async () => await task);
             Assert.True(task.IsCanceled);
             mockedClientWebSocket.Verify(
                 m => m.SendAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(),
                     It.IsAny<CancellationToken>()), Times.Never);
             mockedClientWebSocket.Verify(
                 m => m.CloseAsync(It.IsAny<WebSocketCloseStatus>(), It.IsAny<string>(), It.IsAny<CancellationToken>()),
                 Times.Never);
             mockedClientWebSocket.Verify(
                 m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), token), Times.Never);
         }

         [Fact]
         public async Task ShouldContinueSubmittingOtherMessagesIfOneIsCancelled()
         {
             var mockedClientWebSocket = new Mock<IClientWebSocket>();
             mockedClientWebSocket
                 .SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);
             var tcs = new TaskCompletionSource();
             mockedClientWebSocket.Setup(m => m.SendAsync(It.IsAny<ArraySegment<byte>>(),
                 It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(), It.IsAny<CancellationToken>())).Returns(tcs.Task);
             var connection = GetConnection(mockedClientWebSocket);
             var cts = new CancellationTokenSource();

             var taskToCancel = connection.SubmitAsync<object>(RequestMessage.Build(string.Empty).Create(),
                 cts.Token);
             var taskToComplete = connection.SubmitAsync<object>(RequestMessage.Build(string.Empty).Create(),
                 CancellationToken.None);
             cts.Cancel();
             var taskToComplete2 = connection.SubmitAsync<object>(RequestMessage.Build(string.Empty).Create(),
                 CancellationToken.None);
             tcs.TrySetResult();

             await Assert.ThrowsAsync<TaskCanceledException>(async () => await taskToCancel);
             Assert.True(taskToCancel.IsCanceled);
             await Task.Delay(TimeSpan.FromMilliseconds(200)); // wait a bit to let the messages being sent
             mockedClientWebSocket.Verify(m => m.SendAsync(It.IsAny<ArraySegment<byte>>(),
                 It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()), Times.Exactly(3));
         }

         [Fact]
         public async Task ShouldContinueSubmittingOtherMessagesIfOneIsAlreadyCancelled()
         {
             var mockedClientWebSocket = new Mock<IClientWebSocket>();
             mockedClientWebSocket
                 .SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);
             var cancelledToken = new CancellationToken(canceled: true);
             mockedClientWebSocket
                 .Setup(m => m.SendAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<WebSocketMessageType>(),
                     It.IsAny<bool>(), cancelledToken))
                 .ThrowsAsync(new TaskCanceledException(null, null, cancelledToken));
             var messageToSend = RequestMessage.Build(string.Empty).Create();
             var fakeMessageSerializer = new Mock<IMessageSerializer>();
             var bytesToSend = new byte[] { 1, 2, 3 };
             fakeMessageSerializer.Setup(f => f.SerializeMessageAsync(messageToSend, It.IsAny<CancellationToken>()))
                 .ReturnsAsync(bytesToSend);
             var connection = GetConnection(mockedClientWebSocket, fakeMessageSerializer.Object);

             var taskToCancel = connection.SubmitAsync<object>(RequestMessage.Build(string.Empty).Create(),
                 cancelledToken);
             var taskToComplete = connection.SubmitAsync<object>(messageToSend, CancellationToken.None);

             await Assert.ThrowsAsync<TaskCanceledException>(async () => await taskToCancel);
             Assert.True(taskToCancel.IsCanceled);
             mockedClientWebSocket.Verify(m => m.SendAsync(bytesToSend,
                 It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()), Times.Exactly(1));
         }

         [Fact]
         public async Task ShouldNotProcessReceivedMessageForCancelledRequest()
         {
             var fakeMessageSerializer = new Mock<IMessageSerializer>();
             var receivedBytes = new byte[] { 1, 2, 3 };
             var messageToCancel = RequestMessage.Build(string.Empty).Create();
             var receivedMessage = new ResponseMessage<List<object>>
             {
                 RequestId = messageToCancel.RequestId, Status = new ResponseStatus { Code = ResponseStatusCode.Success }
             };
             fakeMessageSerializer.Setup(f => f.DeserializeMessageAsync(receivedBytes, It.IsAny<CancellationToken>()))
                 .ReturnsAsync(receivedMessage);
             var fakeWebSocketConnection = new Mock<IWebSocketConnection>();
             var receiveTaskCompletionSource = new TaskCompletionSource<byte[]>();
             fakeWebSocketConnection.Setup(m => m.ReceiveMessageAsync()).Returns(receiveTaskCompletionSource.Task);
             var connection = GetConnection(fakeWebSocketConnection.Object, fakeMessageSerializer.Object);
             await connection.ConnectAsync(CancellationToken.None);
             var cts = new CancellationTokenSource();

             var submitTask = connection.SubmitAsync<object>(messageToCancel, cts.Token);
             cts.Cancel();
             receiveTaskCompletionSource.SetResult(receivedBytes);
             await Assert.ThrowsAsync<TaskCanceledException>(() => submitTask);

             Assert.Equal(0, connection.NrRequestsInFlight);
         }

         private static Connection GetConnection(IMock<IClientWebSocket> mockedClientWebSocket,
             IMessageSerializer messageSerializer = null, Uri uri = null)
         {
             return GetConnection(new WebSocketConnection(mockedClientWebSocket.Object, new WebSocketSettings()),
                 messageSerializer, uri);
         }

         private static Connection GetConnection(IWebSocketConnection webSocketConnection,
             IMessageSerializer messageSerializer = null, Uri uri = null)
         {
             uri ??= new Uri("wss://localhost:8182");
             messageSerializer ??= new GraphBinaryMessageSerializer();
             return new Connection(
                 webSocketConnection,
                 uri: uri,
                 username: "user",
                 password: "password",
                 messageSerializer: messageSerializer,
                 sessionId: null);
         }

         private static async Task AssertExpectedConnectionClosedException(WebSocketCloseStatus? expectedCloseStatus, string expectedCloseDescription, Func<Task> func)
         {
             ConnectionClosedException exception = await Assert.ThrowsAsync<ConnectionClosedException>(func);
             Assert.Equal(expectedCloseStatus, exception.Status);
             Assert.Equal(expectedCloseDescription, exception.Description);
         }
     }
 }
	#region License

	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing,
	* software distributed under the License is distributed on an
	* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	* KIND, either express or implied. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*/

	#endregion

	using System;
	using System.Collections.Generic;
	using System.Net.WebSockets;
	using System.Threading;
	using System.Threading.Tasks;
	using Gremlin.Net.Driver;
	using Gremlin.Net.Driver.Exceptions;
	using Gremlin.Net.Driver.Messages;
	using Gremlin.Net.Structure.IO.GraphBinary;
	using Moq;
	using Xunit;

	namespace Gremlin.Net.UnitTest.Driver
	{
	public class ConnectionTests
	{
	[Fact]
	public async Task ShouldHandleCloseMessageAfterConnectAsync()
	{
	var mockedClientWebSocket = new Mock<IClientWebSocket>();
	mockedClientWebSocket
	.Setup(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()))
	.ReturnsAsync(new WebSocketReceiveResult(0, WebSocketMessageType.Close, true, WebSocketCloseStatus.MessageTooBig, "Message is too large"));
	mockedClientWebSocket
	.SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);

	Uri uri = new Uri("wss://localhost:8182");
	Connection connection = GetConnection(mockedClientWebSocket, uri: uri);

	await connection.ConnectAsync(CancellationToken.None);

	Assert.False(connection.IsOpen);
	Assert.Equal(0, connection.NrRequestsInFlight);
	mockedClientWebSocket.Verify(m => m.ConnectAsync(uri, It.IsAny<CancellationToken>()), Times.Once);
	mockedClientWebSocket.Verify(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()), Times.Once);
	mockedClientWebSocket.Verify(m => m.CloseAsync(WebSocketCloseStatus.NormalClosure, string.Empty, It.IsAny<CancellationToken>()), Times.Once);
	}

	[Fact]
	public async Task ShouldThrowIfClosedMessageReceivedWithValidPropertiesAsync()
	{
	var mockedClientWebSocket = new Mock<IClientWebSocket>();
	mockedClientWebSocket
	.SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);

	WebSocketConnection webSocketConnection = new WebSocketConnection(
	mockedClientWebSocket.Object,
	new WebSocketSettings());

	// Test all known close statuses
	foreach (Enum closeStatus in Enum.GetValues(typeof(WebSocketCloseStatus)))
	{
	mockedClientWebSocket
	.Setup(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()))
	.ReturnsAsync(new WebSocketReceiveResult(0, WebSocketMessageType.Close, true, (WebSocketCloseStatus)closeStatus, closeStatus.ToString()));

	await AssertExpectedConnectionClosedException((WebSocketCloseStatus?)closeStatus, closeStatus.ToString(), () => webSocketConnection.ReceiveMessageAsync());
	}

	// Test null/empty close property values as well.
	mockedClientWebSocket
	.Setup(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()))
	.ReturnsAsync(new WebSocketReceiveResult(0, WebSocketMessageType.Close, true, null, null));
	await AssertExpectedConnectionClosedException(null, null, () => webSocketConnection.ReceiveMessageAsync());

	mockedClientWebSocket
	.Setup(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()))
	.ReturnsAsync(new WebSocketReceiveResult(0, WebSocketMessageType.Close, true, null, String.Empty));
	await AssertExpectedConnectionClosedException(null, String.Empty, () => webSocketConnection.ReceiveMessageAsync());
	}

	[Fact]
	public async Task ShouldThrowOnSubmitRequestIfWebSocketIsNotOpen()
	{
	// This is to test that race bugs don't get introduced which would cause submitted requests to hang if the
	// connection is being closed/aborted or is not connected. In these cases, WebSocket.SendAsync should throw for the underlying
	// websocket is not open and the caller should be notified of that failure.
	var mockedClientWebSocket = new Mock<IClientWebSocket>();

	mockedClientWebSocket
	.SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);

	Connection connection = GetConnection(mockedClientWebSocket);
	RequestMessage request = RequestMessage.Build("gremlin").Create();

	// Simulate the SendAsync exception behavior if the underlying websocket is closed (see reference https://docs.microsoft.com/en-us/dotnet/api/system.net.websockets.clientwebsocket.sendasync?view=net-6.0)
	mockedClientWebSocket.Setup(m => m.SendAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()))
	.ThrowsAsync(new ObjectDisposedException(nameof(ClientWebSocket), "Socket closed"));

	// Test various closing/closed WebSocketStates with SubmitAsync.
	mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.Closed);
	await Assert.ThrowsAsync<ObjectDisposedException>(() => connection.SubmitAsync<dynamic>(request, CancellationToken.None));

	mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.CloseSent);
	await Assert.ThrowsAsync<ObjectDisposedException>(() => connection.SubmitAsync<dynamic>(request, CancellationToken.None));

	mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.CloseReceived);
	await Assert.ThrowsAsync<ObjectDisposedException>(() => connection.SubmitAsync<dynamic>(request, CancellationToken.None));

	mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.Aborted);
	await Assert.ThrowsAsync<ObjectDisposedException>(() => connection.SubmitAsync<dynamic>(request, CancellationToken.None));

	// Simulate SendAsync exception behavior if underlying websocket is not connected.
	mockedClientWebSocket.Setup(m => m.SendAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()))
	.ThrowsAsync(new InvalidOperationException("Socket not connected"));
	mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.Connecting);
	await Assert.ThrowsAsync<InvalidOperationException>(() => connection.SubmitAsync<dynamic>(request, CancellationToken.None));
	}

	[Fact]
	public async Task ShouldHandleCloseMessageForInFlightRequestsAsync()
	{
	// Tests that in-flight requests will get notified if a connection close message is received.
	Uri uri = new Uri("wss://localhost:8182");
	WebSocketReceiveResult closeResult = new WebSocketReceiveResult(0, WebSocketMessageType.Close, true, WebSocketCloseStatus.EndpointUnavailable, "Server shutdown");

	var receiveSempahore = new SemaphoreSlim(0, 1);
	var mockedClientWebSocket = new Mock<IClientWebSocket>();
	mockedClientWebSocket
	.Setup(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()))
	.Returns(async () =>
	{
	await receiveSempahore.WaitAsync();
	mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.CloseReceived);
	return closeResult;
	});
	mockedClientWebSocket.Setup(m => m.State).Returns(WebSocketState.Open);
	mockedClientWebSocket
	.SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);

	var connection = GetConnection(mockedClientWebSocket, uri: uri);
	await connection.ConnectAsync(CancellationToken.None);

	// Create two in-flight requests that will block on waiting for a response.
	RequestMessage requestMsg1 = RequestMessage.Build("gremlin").Create();
	RequestMessage requestMsg2 = RequestMessage.Build("gremlin").Create();
	Task request1 = connection.SubmitAsync<dynamic>(requestMsg1, CancellationToken.None);
	Task request2 = connection.SubmitAsync<dynamic>(requestMsg2, CancellationToken.None);

	// Confirm the requests are in-flight.
	Assert.Equal(2, connection.NrRequestsInFlight);

	// Release the connection close message.
	receiveSempahore.Release();

	// Assert that both requests get notified with the closed exception.
	await AssertExpectedConnectionClosedException(closeResult.CloseStatus, closeResult.CloseStatusDescription, () => request1);
	await AssertExpectedConnectionClosedException(closeResult.CloseStatus, closeResult.CloseStatusDescription, () => request2);

	// delay for NotifyAboutConnectionFailure running in another thread
	var runs = 0;
	while (connection.NrRequestsInFlight == 2 && ++runs < 5)
	{
	await Task.Delay(5);
	}

	Assert.False(connection.IsOpen);
	Assert.Equal(0, connection.NrRequestsInFlight);
	mockedClientWebSocket.Verify(m => m.ConnectAsync(uri, It.IsAny<CancellationToken>()), Times.Once);
	mockedClientWebSocket.Verify(m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<CancellationToken>()), Times.Once);
	mockedClientWebSocket.Verify(m => m.CloseAsync(WebSocketCloseStatus.NormalClosure, string.Empty, It.IsAny<CancellationToken>()), Times.Once);
	}

	[Fact]
	public async Task ShouldProperlyHandleCancellationForSubmitAsync()
	{
	var mockedClientWebSocket = new Mock<IClientWebSocket>();
	mockedClientWebSocket
	.SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);
	var connection = GetConnection(mockedClientWebSocket);
	var cts = new CancellationTokenSource();

	var task = connection.SubmitAsync<object>(RequestMessage.Build(string.Empty).Create(),
	cts.Token);
	cts.Cancel();

	await Assert.ThrowsAsync<TaskCanceledException>(async () => await task);
	Assert.True(task.IsCanceled);
	mockedClientWebSocket.Verify(m => m.SendAsync(It.IsAny<ArraySegment<byte>>(),
	It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(), cts.Token), Times.Once);
	mockedClientWebSocket.Verify(
	m => m.CloseAsync(It.IsAny<WebSocketCloseStatus>(), It.IsAny<string>(), It.IsAny<CancellationToken>()),
	Times.Never);
	mockedClientWebSocket.Verify(
	m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), cts.Token), Times.Never);
	}

	[Fact]
	public async Task ShouldProperlyHandleCancellationForSubmitAsyncIfAlreadyCancelled()
	{
	var mockedClientWebSocket = new Mock<IClientWebSocket>();
	mockedClientWebSocket
	.SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);
	var connection = GetConnection(mockedClientWebSocket);
	var token = new CancellationToken(canceled: true);

	var task = connection.SubmitAsync<object>(RequestMessage.Build(string.Empty).Create(),
	token);

	await Assert.ThrowsAsync<TaskCanceledException>(async () => await task);
	Assert.True(task.IsCanceled);
	mockedClientWebSocket.Verify(
	m => m.SendAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(),
	It.IsAny<CancellationToken>()), Times.Never);
	mockedClientWebSocket.Verify(
	m => m.CloseAsync(It.IsAny<WebSocketCloseStatus>(), It.IsAny<string>(), It.IsAny<CancellationToken>()),
	Times.Never);
	mockedClientWebSocket.Verify(
	m => m.ReceiveAsync(It.IsAny<ArraySegment<byte>>(), token), Times.Never);
	}

	[Fact]
	public async Task ShouldContinueSubmittingOtherMessagesIfOneIsCancelled()
	{
	var mockedClientWebSocket = new Mock<IClientWebSocket>();
	mockedClientWebSocket
	.SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);
	var tcs = new TaskCompletionSource();
	mockedClientWebSocket.Setup(m => m.SendAsync(It.IsAny<ArraySegment<byte>>(),
	It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(), It.IsAny<CancellationToken>())).Returns(tcs.Task);
	var connection = GetConnection(mockedClientWebSocket);
	var cts = new CancellationTokenSource();

	var taskToCancel = connection.SubmitAsync<object>(RequestMessage.Build(string.Empty).Create(),
	cts.Token);
	var taskToComplete = connection.SubmitAsync<object>(RequestMessage.Build(string.Empty).Create(),
	CancellationToken.None);
	cts.Cancel();
	var taskToComplete2 = connection.SubmitAsync<object>(RequestMessage.Build(string.Empty).Create(),
	CancellationToken.None);
	tcs.TrySetResult();

	await Assert.ThrowsAsync<TaskCanceledException>(async () => await taskToCancel);
	Assert.True(taskToCancel.IsCanceled);
	await Task.Delay(TimeSpan.FromMilliseconds(200)); // wait a bit to let the messages being sent
	mockedClientWebSocket.Verify(m => m.SendAsync(It.IsAny<ArraySegment<byte>>(),
	It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()), Times.Exactly(3));
	}

	[Fact]
	public async Task ShouldContinueSubmittingOtherMessagesIfOneIsAlreadyCancelled()
	{
	var mockedClientWebSocket = new Mock<IClientWebSocket>();
	mockedClientWebSocket
	.SetupGet(m => m.Options).Returns(new ClientWebSocket().Options);
	var cancelledToken = new CancellationToken(canceled: true);
	mockedClientWebSocket
	.Setup(m => m.SendAsync(It.IsAny<ArraySegment<byte>>(), It.IsAny<WebSocketMessageType>(),
	It.IsAny<bool>(), cancelledToken))
	.ThrowsAsync(new TaskCanceledException(null, null, cancelledToken));
	var messageToSend = RequestMessage.Build(string.Empty).Create();
	var fakeMessageSerializer = new Mock<IMessageSerializer>();
	var bytesToSend = new byte[] { 1, 2, 3 };
	fakeMessageSerializer.Setup(f => f.SerializeMessageAsync(messageToSend, It.IsAny<CancellationToken>()))
	.ReturnsAsync(bytesToSend);
	var connection = GetConnection(mockedClientWebSocket, fakeMessageSerializer.Object);

	var taskToCancel = connection.SubmitAsync<object>(RequestMessage.Build(string.Empty).Create(),
	cancelledToken);
	var taskToComplete = connection.SubmitAsync<object>(messageToSend, CancellationToken.None);

	await Assert.ThrowsAsync<TaskCanceledException>(async () => await taskToCancel);
	Assert.True(taskToCancel.IsCanceled);
	mockedClientWebSocket.Verify(m => m.SendAsync(bytesToSend,
	It.IsAny<WebSocketMessageType>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()), Times.Exactly(1));
	}

	[Fact]
	public async Task ShouldNotProcessReceivedMessageForCancelledRequest()
	{
	var fakeMessageSerializer = new Mock<IMessageSerializer>();
	var receivedBytes = new byte[] { 1, 2, 3 };
	var messageToCancel = RequestMessage.Build(string.Empty).Create();
	var receivedMessage = new ResponseMessage<List<object>>
	{
	RequestId = messageToCancel.RequestId, Status = new ResponseStatus { Code = ResponseStatusCode.Success }
	};
	fakeMessageSerializer.Setup(f => f.DeserializeMessageAsync(receivedBytes, It.IsAny<CancellationToken>()))
	.ReturnsAsync(receivedMessage);
	var fakeWebSocketConnection = new Mock<IWebSocketConnection>();
	var receiveTaskCompletionSource = new TaskCompletionSource<byte[]>();
	fakeWebSocketConnection.Setup(m => m.ReceiveMessageAsync()).Returns(receiveTaskCompletionSource.Task);
	var connection = GetConnection(fakeWebSocketConnection.Object, fakeMessageSerializer.Object);
	await connection.ConnectAsync(CancellationToken.None);
	var cts = new CancellationTokenSource();

	var submitTask = connection.SubmitAsync<object>(messageToCancel, cts.Token);
	cts.Cancel();
	receiveTaskCompletionSource.SetResult(receivedBytes);
	await Assert.ThrowsAsync<TaskCanceledException>(() => submitTask);

	Assert.Equal(0, connection.NrRequestsInFlight);
	}

	private static Connection GetConnection(IMock<IClientWebSocket> mockedClientWebSocket,
	IMessageSerializer messageSerializer = null, Uri uri = null)
	{
	return GetConnection(new WebSocketConnection(mockedClientWebSocket.Object, new WebSocketSettings()),
	messageSerializer, uri);
	}

	private static Connection GetConnection(IWebSocketConnection webSocketConnection,
	IMessageSerializer messageSerializer = null, Uri uri = null)
	{
	uri ??= new Uri("wss://localhost:8182");
	messageSerializer ??= new GraphBinaryMessageSerializer();
	return new Connection(
	webSocketConnection,
	uri: uri,
	username: "user",
	password: "password",
	messageSerializer: messageSerializer,
	sessionId: null);
	}

	private static async Task AssertExpectedConnectionClosedException(WebSocketCloseStatus? expectedCloseStatus, string expectedCloseDescription, Func<Task> func)
	{
	ConnectionClosedException exception = await Assert.ThrowsAsync<ConnectionClosedException>(func);
	Assert.Equal(expectedCloseStatus, exception.Status);
	Assert.Equal(expectedCloseDescription, exception.Description);
	}
	}
	}