From cb0236f592e1217d202184356d8b4a96f6703534 Mon Sep 17 00:00:00 2001
From: Turo Lamminen <turo.lamminen@alternativegames.net>
Date: Thu, 9 Dec 2021 16:48:59 +0200
Subject: [PATCH] Add IWebSocket interface

---
 Runtime/SimpleGraphQL/GraphQLClient.cs | 10 +--
 Runtime/SimpleGraphQL/HttpUtils.cs     | 97 +++++++++++++++++++++++++-
 2 files changed, 100 insertions(+), 7 deletions(-)

diff --git a/Runtime/SimpleGraphQL/GraphQLClient.cs b/Runtime/SimpleGraphQL/GraphQLClient.cs
index 0359af2..dcdce8a 100644
--- a/Runtime/SimpleGraphQL/GraphQLClient.cs
+++ b/Runtime/SimpleGraphQL/GraphQLClient.cs
@@ -134,10 +134,11 @@ public async Task<bool> SubscribeAsync(
             string id,
             string authToken = null,
             Dictionary<string, object> variables = null,
-            Dictionary<string, string> headers = null
+            Dictionary<string, string> headers = null,
+            HttpUtils.IWebSocket ws = null
         )
         {
-            return await SubscribeAsync(query, id, AuthScheme, authToken, variables, headers);
+            return await SubscribeAsync(query, id, AuthScheme, authToken, variables, headers, ws);
         }
 
         public async Task<bool> SubscribeAsync(
@@ -146,7 +147,8 @@ public async Task<bool> SubscribeAsync(
             string authScheme = "Bearer",
             string authToken = null,
             Dictionary<string, object> variables = null,
-            Dictionary<string, string> headers = null
+            Dictionary<string, string> headers = null,
+            HttpUtils.IWebSocket ws = null
         )
         {
             if(query.OperationType != OperationType.Subscription)
@@ -171,7 +173,7 @@ public async Task<bool> SubscribeAsync(
             if(!HttpUtils.IsWebSocketReady())
             {
                 // Prepare the socket before continuing.
-                bool connectionSuccessful = await HttpUtils.WebSocketConnect(Endpoint, authScheme, authToken, "graphql-ws", headers);
+                bool connectionSuccessful = await HttpUtils.WebSocketConnect(Endpoint, authScheme, authToken, "graphql-ws", headers, ws);
                 if(!connectionSuccessful)
                 {
                     return false;
diff --git a/Runtime/SimpleGraphQL/HttpUtils.cs b/Runtime/SimpleGraphQL/HttpUtils.cs
index 1acb0ef..7855e30 100644
--- a/Runtime/SimpleGraphQL/HttpUtils.cs
+++ b/Runtime/SimpleGraphQL/HttpUtils.cs
@@ -21,7 +21,90 @@ public enum SubscriptionError
     [PublicAPI]
     public static class HttpUtils
     {
-        private static ClientWebSocket _webSocket;
+        public interface IWebSocket : IDisposable
+        {
+            interface ISocketOptions
+            {
+                public void AddSubProtocol(string protocol);
+
+                public void SetRequestHeader(string v1, string v2);
+            }
+
+            WebSocketState State { get; }
+            ISocketOptions Options { get; }
+            WebSocketCloseStatus? CloseStatus { get; }
+
+            Task ConnectAsync(Uri uri, CancellationToken none);
+            Task SendAsync(ArraySegment<byte> arraySegment, WebSocketMessageType text, bool v, CancellationToken ct);
+            Task CloseAsync(WebSocketCloseStatus normalClosure, string v, CancellationToken ct);
+            Task<WebSocketReceiveResult> ReceiveAsync(ArraySegment<byte> buffer, CancellationToken ct);
+        }
+
+
+        public class ClientWebSocketOptionsWrapper : IWebSocket.ISocketOptions
+        {
+            ClientWebSocketOptions options;
+
+            public ClientWebSocketOptionsWrapper(ClientWebSocketOptions options)
+            {
+                this.options = options;
+            }
+
+            public void AddSubProtocol(string protocol)
+            {
+                options.AddSubProtocol(protocol);
+            }
+
+            public void SetRequestHeader(string v1, string v2)
+            {
+                options.SetRequestHeader(v1, v2);
+            }
+        }
+
+
+        public class ClientWebSocketWrapper : IWebSocket
+        {
+            ClientWebSocket ws;
+
+            public ClientWebSocketWrapper()
+            {
+                ws = new ClientWebSocket();
+            }
+
+            WebSocketState IWebSocket.State => ws.State;
+
+            public WebSocketCloseStatus? CloseStatus => ws.CloseStatus;
+
+            public IWebSocket.ISocketOptions Options => new ClientWebSocketOptionsWrapper(ws.Options);
+
+
+            public Task ConnectAsync(Uri uri, CancellationToken ct)
+            {
+                return ws.ConnectAsync(uri, ct);
+            }
+
+            public Task SendAsync(ArraySegment<byte> arraySegment, WebSocketMessageType text, bool v, CancellationToken ct)
+            {
+                return ws.SendAsync(arraySegment, text, v, ct);
+            }
+
+            public Task CloseAsync(WebSocketCloseStatus normalClosure, string v, CancellationToken ct)
+            {
+                return ws.CloseAsync(normalClosure, v, ct);
+            }
+
+            public Task<WebSocketReceiveResult> ReceiveAsync(ArraySegment<byte> buffer, CancellationToken ct)
+            {
+                return ws.ReceiveAsync(buffer, ct);
+            }
+
+            public void Dispose()
+            {
+                ws.Dispose();
+            }
+        }
+
+        private static IWebSocket _webSocket;
 
         /// <summary>
         /// Called when the websocket receives subscription data.
@@ -119,13 +202,21 @@ public static async Task<bool> WebSocketConnect(
             string authScheme = "Bearer",
             string authToken = null,
             string protocol = "graphql-ws",
-            Dictionary<string, string> headers = null
+            Dictionary<string, string> headers = null,
+            IWebSocket ws = null
         )
         {
             url = url.Replace("http", "ws");
 
             Uri uri = new Uri(url);
-            _webSocket = new ClientWebSocket();
+            if (ws != null)
+            {
+                _webSocket = ws;
+            }
+            else
+            {
+                _webSocket = new ClientWebSocketWrapper();
+            }
             _webSocket.Options.AddSubProtocol(protocol);
 
             if(authToken != null)