From aa9e71b7e0d8c6f5a6dd155b89ba3a1764a95698 Mon Sep 17 00:00:00 2001
From: Benjamin Himes <37844818+thewhaleking@users.noreply.github.com>
Date: Wed, 27 Nov 2024 12:15:23 +0200
Subject: [PATCH] Update `ensure_connected` for websockets (#2486)

* Updates `ensure_connected` for the change to websockets.

* Regenerates the websocket connection if it is closed.
---
 bittensor/core/subtensor.py   | 2 +-
 bittensor/utils/networking.py | 4 ++--
 tests/helpers/helpers.py      | 2 ++
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/bittensor/core/subtensor.py b/bittensor/core/subtensor.py
index 6c17cc87b0..fa46242277 100644
--- a/bittensor/core/subtensor.py
+++ b/bittensor/core/subtensor.py
@@ -216,7 +216,7 @@ def _get_substrate(self):
         """Establishes a connection to the Substrate node using configured parameters."""
         try:
             # Set up params.
-            if not self.websocket:
+            if self.websocket is None or self.websocket.close_code is not None:
                 self.websocket = ws_client.connect(
                     self.chain_endpoint,
                     open_timeout=self._connection_timeout,
diff --git a/bittensor/utils/networking.py b/bittensor/utils/networking.py
index aca8ee2b04..f47c88512e 100644
--- a/bittensor/utils/networking.py
+++ b/bittensor/utils/networking.py
@@ -10,7 +10,7 @@
 import netaddr
 import requests
 from retry import retry
-from websocket import WebSocketConnectionClosedException
+from websockets.exceptions import ConnectionClosed
 
 from bittensor.utils.btlogging import logging
 
@@ -193,7 +193,7 @@ def wrapper(self, *args, **kwargs):
 
         try:
             return func(self, *args, **kwargs)
-        except WebSocketConnectionClosedException:
+        except ConnectionClosed:
             logging.console.warning(
                 "WebSocket connection closed. Attempting to reconnect 5 times..."
             )
diff --git a/tests/helpers/helpers.py b/tests/helpers/helpers.py
index 1a7581ccd5..a125e15329 100644
--- a/tests/helpers/helpers.py
+++ b/tests/helpers/helpers.py
@@ -125,6 +125,8 @@ def get_mock_neuron_by_uid(uid: int, **kwargs) -> NeuronInfo:
 
 
 class FakeWebsocket(ClientConnection):
+    close_code = None
+
     def __init__(self, *args, seed, **kwargs):
         protocol = ClientProtocol(parse_uri("ws://127.0.0.1:9945"))
         super().__init__(socket=None, protocol=protocol, **kwargs)