From 604048e346ac0990cb970c1dd75ed3a10bc77cb0 Mon Sep 17 00:00:00 2001
From: Apolixit <friotromain@gmail.com>
Date: Sat, 14 Sep 2024 12:17:36 +0200
Subject: [PATCH] Add manually disconnect + Minor fix post merge

---
 Substrate.NetApi.TestNode/ClientTests.cs | 22 +++++++++++
 Substrate.NetApi/SubstrateClient.cs      | 48 +++++++++++++++++++++++-
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/Substrate.NetApi.TestNode/ClientTests.cs b/Substrate.NetApi.TestNode/ClientTests.cs
index 1f64780..4bb339a 100644
--- a/Substrate.NetApi.TestNode/ClientTests.cs
+++ b/Substrate.NetApi.TestNode/ClientTests.cs
@@ -70,5 +70,27 @@ public async Task OnConnectionLost_ShouldThrowDisconnectedEventAsync()
             await Task.WhenAny(onConnectionLostTriggered.Task, Task.Delay(TimeSpan.FromMinutes(1)));
             Assert.That(onConnectionLostTriggered.Task.IsCompleted, Is.True);
         }
+
+        [Test]
+        public async Task ManuallyDisconnect_ShouldNotTryToReconnectAsync()
+        {
+            await _client.ConnectAsync();
+            await _client.CloseAsync();
+
+            Assert.That(_client.IsConnected, Is.False);
+        }
+
+        [Test]
+        public async Task Disconnect_ShouldTryToReconnectAsync()
+        {
+            var onReconnectedTriggered = new TaskCompletionSource<(bool, int)>();
+            _client.OnReconnected += (sender, e) => onReconnectedTriggered.SetResult((true, e));
+
+            await _client.ConnectAsync();
+            await _client._socket.CloseAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None);
+
+            await Task.WhenAny(onReconnectedTriggered.Task, Task.Delay(TimeSpan.FromMinutes(1)));
+            Assert.That(_client.IsConnected, Is.True);
+        }
     }
 }
diff --git a/Substrate.NetApi/SubstrateClient.cs b/Substrate.NetApi/SubstrateClient.cs
index 0333514..82745d8 100644
--- a/Substrate.NetApi/SubstrateClient.cs
+++ b/Substrate.NetApi/SubstrateClient.cs
@@ -23,6 +23,7 @@
 using Substrate.NetApi.Model.Types.Metadata.V14;
 
 [assembly: InternalsVisibleTo("Substrate.NetApi.Test")]
+[assembly: InternalsVisibleTo("Substrate.NetApi.TestNode")]
 
 namespace Substrate.NetApi
 {
@@ -63,7 +64,12 @@ public class SubstrateClient : IDisposable
         private JsonRpc _jsonRpc;
 
         /// <summary> The socket. </summary>
-        private ClientWebSocket _socket;
+        internal ClientWebSocket _socket;
+
+        /// <summary>
+        /// Check if the client has been disconnected manually (to avoid auto-reconnect)
+        /// </summary>
+        private bool _isDisconnectedManually = false;
 
         /// <summary>
         /// The "ping" to check the connection status
@@ -80,6 +86,12 @@ public class SubstrateClient : IDisposable
         /// </summary>
         public event EventHandler ConnectionSet;
 
+        /// <summary>
+        /// Event triggered when the connection is reconnected
+        /// </summary>
+
+        public event EventHandler<int> OnReconnected;
+
         /// <summary>
         /// Bypass Remote Certificate Validation. Useful when testing with self-signed SSL certificates.
         /// </summary>
@@ -198,6 +210,30 @@ public bool SetJsonRPCTraceLevel(SourceLevels sourceLevels)
             return true;
         }
 
+        /// <summary>
+        /// Raises the event when the connection to the server is lost.
+        /// </summary>
+        protected virtual void OnConnectionLost()
+        {
+            ConnectionLost?.Invoke(this, EventArgs.Empty);
+        }
+
+        /// <summary>
+        /// Raises the event when the connection to the server is set.
+        /// </summary>
+        protected virtual void OnConnectionSet()
+        {
+            ConnectionSet?.Invoke(this, EventArgs.Empty);
+        }
+
+        /// <summary>
+        /// Raises the event when reconnected
+        /// </summary>
+        protected virtual void OnReconnectedSet(int nbTry)
+        {
+            OnReconnected?.Invoke(this, nbTry);
+        }
+
         /// <summary>
         /// Asynchronously connects to the node.
         /// </summary>
@@ -350,6 +386,13 @@ public async Task ConnectAsync(bool useMetaData, bool standardSubstrate, int max
         private void OnJsonRpcDisconnected(object sender, JsonRpcDisconnectedEventArgs e)
         {
             Logger.Error(e.Exception, $"JsonRpc disconnected: {e.Reason}");
+            OnConnectionLost();
+
+            if(_isDisconnectedManually)
+            {
+                _isDisconnectedManually = false;
+                return;
+            }
 
             // Attempt to reconnect asynchronously
             _ = Task.Run(async () =>
@@ -385,6 +428,8 @@ await ConnectAsync(
                     );
 
                     Logger.Information("Reconnected successfully.");
+
+                    OnReconnectedSet(retry);
                 }
                 catch (Exception ex)
                 {
@@ -579,6 +624,7 @@ public async Task CloseAsync()
         public async Task CloseAsync(CancellationToken token)
         {
             _connectTokenSource?.Cancel();
+            _isDisconnectedManually = true;
 
             await Task.Run(async () =>
             {