MockSubtensor work offline (#2487)

* Updates `ensure_connected` for the change to websockets.

* Only need to catch ConnectionClosed, as ConnectionClosedError subclasses ConnectionClosed.

* Regenerates the websocket connection if it is closed.

* Add property to FakeWebsocket

* Ruff

* Makes tests that use MockSubtensor run offline.

* Changed scope of patch.

Author: thewhaleking

bittensor/utils/mock/subtensor_mock.py

@@ -20,9 +20,11 @@
 from hashlib import sha256
 from types import SimpleNamespace
 from typing import Any, Optional, Union, TypedDict
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 from bittensor_wallet import Wallet
+from substrateinterface.base import SubstrateInterface
+from websockets.sync.client import ClientConnection
 
 from bittensor.core.chain_data import (
     NeuronInfo,
@@ -33,6 +35,7 @@
 from bittensor.core.types import AxonServeCallParams, PrometheusServeCallParams
 from bittensor.core.errors import ChainQueryError
 from bittensor.core.subtensor import Subtensor
+import bittensor.core.subtensor as subtensor_module
 from bittensor.utils import RAOPERTAO, u16_normalized_float
 from bittensor.utils.balance import Balance
 
@@ -248,14 +251,22 @@ def setup(self) -> None:
 
             self.network = "mock"
             self.chain_endpoint = "ws://mock_endpoint.bt"
-            self.substrate = MagicMock()
+            self.substrate = MagicMock(autospec=SubstrateInterface)
 
     def __init__(self, *args, **kwargs) -> None:
-        super().__init__()
-        self.__dict__ = __GLOBAL_MOCK_STATE__
-
-        if not hasattr(self, "chain_state") or getattr(self, "chain_state") is None:
-            self.setup()
+        mock_substrate_interface = MagicMock(autospec=SubstrateInterface)
+        mock_websocket = MagicMock(autospec=ClientConnection)
+        mock_websocket.close_code = None
+        with patch.object(
+            subtensor_module,
+            "SubstrateInterface",
+            return_value=mock_substrate_interface,
+        ):
+            super().__init__(websocket=mock_websocket)
+            self.__dict__ = __GLOBAL_MOCK_STATE__
+
+            if not hasattr(self, "chain_state") or getattr(self, "chain_state") is None:
+                self.setup()
 
     def get_block_hash(self, block_id: int) -> str:
         return "0x" + sha256(str(block_id).encode()).hexdigest()[:64]