clean up

src/NBomber.WebSockets/WebSocket.cs

@@ -1,4 +1,5 @@
 using System.Net.WebSockets;
+using System.Text;
 using Microsoft.IO;
 
 namespace NBomber.WebSockets;
@@ -18,6 +19,62 @@ public class WebSocket(WebSocketConfig config) : IDisposable
 
     public ClientWebSocket Client { get; } = new();
 
+    /// <summary>
+    /// This method should be used to connect to a WebSocket server asynchronously.
+    /// </summary>
+    /// <param name="url">The URL of the WebSocket server to connect to.</param>
+    /// <param name="cancellationToken">A <see cref="CancellationToken"/> used to propagate notification that the operation should be canceled.</param>
+    /// <exception cref="WebSocketException">Thrown when an error occurs during WebSocket communication.</exception>
+    /// <exception cref="OperationCanceledException">Thrown if the receive operation is canceled by the provided cancellation token.</exception>
+    public Task Connect(string url, CancellationToken cancellationToken = default)
+    {
+        return Client.ConnectAsync(new Uri(url), cancellationToken);
+    }
+
+    /// <summary>
+    /// This method should be used to connect to a WebSocket server asynchronously.
+    /// </summary>
+    /// <param name="uri">The URI of the WebSocket server to connect to.</param>
+    /// <param name="cancellationToken">A <see cref="CancellationToken"/> used to propagate notification that the operation should be canceled.</param>
+    /// <exception cref="WebSocketException">Thrown when an error occurs during WebSocket communication.</exception>
+    /// <exception cref="OperationCanceledException">Thrown if the receive operation is canceled by the provided cancellation token.</exception>
+    public Task Connect(Uri uri, CancellationToken cancellationToken = default)
+    {
+        return Client.ConnectAsync(uri, cancellationToken);
+    }
+
+    /// <summary>
+    /// This method should be used to send data on WebSocket asynchronously.
+    /// </summary>
+    /// <param name="text">The text to be sent. It will be encoded using UTF8 encoding.</param>
+    /// <param name="cancellationToken">A <see cref="CancellationToken"/> used to propagate notification that the operation should be canceled.</param>
+    /// <exception cref="WebSocketException">Thrown when an error occurs during WebSocket communication.</exception>
+    /// <exception cref="OperationCanceledException">Thrown if the receive operation is canceled by the provided cancellation token.</exception>
+    public ValueTask Send(string text, CancellationToken cancellationToken = default)
+    {
+        using var ms = MsStreamManager.GetStream();
+
+        var buffer = ms.GetMemory(text.Length);
+        var byteCount = Encoding.UTF8.GetBytes(text.AsSpan(), buffer.Span);
+        ms.Advance(byteCount);
+
+        var msg = ms.GetBuffer().AsMemory(0, (int) ms.Length);
+
+        return Client.SendAsync(msg, WebSocketMessageType.Text, true, cancellationToken);
+    }
+
+    /// <summary>
+    /// This method should be used to send data on WebSocket asynchronously.
+    /// </summary>
+    /// <param name="payload">The binary payload to be sent.</param>
+    /// <param name="cancellationToken">A <see cref="CancellationToken"/> used to propagate notification that the operation should be canceled.</param>
+    /// <exception cref="WebSocketException">Thrown when an error occurs during WebSocket communication.</exception>
+    /// <exception cref="OperationCanceledException">Thrown if the receive operation is canceled by the provided cancellation token.</exception>
+    public ValueTask Send(ReadOnlyMemory<byte> payload, CancellationToken cancellationToken = default)
+    {
+        return Client.SendAsync(payload, WebSocketMessageType.Binary, true, cancellationToken);
+    }
+
     /// <summary>
     /// This method should be used to receive a WebSocket message from the connected WebSocket asynchronously. 
     /// This method returns <see cref="WebSocketResponse"/> that should be disposed of after usage.
@@ -57,6 +114,18 @@ public async ValueTask<WebSocketResponse> Receive(CancellationToken cancellation
             throw;
         }
     }
+
+    /// <summary>
+    /// This method should be used to close WebSocket connection asynchronously.
+    /// </summary>
+    /// <param name="closeStatus">The WebSocket close status.</param>
+    /// <param name="cancellationToken">A <see cref="CancellationToken"/> used to propagate notification that the operation should be canceled.</param>
+    /// <exception cref="WebSocketException">Thrown when an error occurs during WebSocket communication.</exception>
+    /// <exception cref="OperationCanceledException">Thrown if the receive operation is canceled by the provided cancellation token.</exception>
+    public Task Close(WebSocketCloseStatus closeStatus = WebSocketCloseStatus.NormalClosure, CancellationToken cancellationToken = default)
+    {
+        return Client.CloseAsync(closeStatus, null, cancellationToken);
+    }
 
     public void Dispose()
     {