Update ensure_connected for websockets

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,

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..."
             )

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)