From 785068669385ba18c6b9161dd82cc54bdf6f125d Mon Sep 17 00:00:00 2001
From: Benjamin Himes <benhimes@opentensor.dev>
Date: Mon, 25 Nov 2024 16:23:53 +0200
Subject: [PATCH] Moved FakeWebsocket to a helper

---
 tests/helpers/helpers.py                      | 35 +++++++++++++++++
 .../test_subtensor_integration.py             | 38 +------------------
 2 files changed, 36 insertions(+), 37 deletions(-)

diff --git a/tests/helpers/helpers.py b/tests/helpers/helpers.py
index 41109ee5e6..1a7581ccd5 100644
--- a/tests/helpers/helpers.py
+++ b/tests/helpers/helpers.py
@@ -15,8 +15,13 @@
 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 # DEALINGS IN THE SOFTWARE.
 
+from collections import deque
+import json
 from typing import Union
 
+from websockets.sync.client import ClientConnection, ClientProtocol
+from websockets.uri import parse_uri
+
 from bittensor_wallet.mock.wallet_mock import MockWallet as _MockWallet
 from bittensor_wallet.mock.wallet_mock import get_mock_coldkey
 from bittensor_wallet.mock.wallet_mock import get_mock_hotkey
@@ -24,6 +29,7 @@
 
 from bittensor.utils.balance import Balance
 from bittensor.core.chain_data import AxonInfo, NeuronInfo, PrometheusInfo
+from tests.helpers.integration_websocket_data import WEBSOCKET_RESPONSES
 
 
 def __mock_wallet_factory__(*_, **__) -> _MockWallet:
@@ -60,6 +66,7 @@ def get_mock_neuron(**kwargs) -> NeuronInfo:
     """
 
     mock_neuron_d = dict(
+        # TODO fix the AxonInfo here — it doesn't work
         {
             "netuid": -1,  # mock netuid
             "axon_info": AxonInfo(
@@ -115,3 +122,31 @@ def get_mock_neuron_by_uid(uid: int, **kwargs) -> NeuronInfo:
     return get_mock_neuron(
         uid=uid, hotkey=get_mock_hotkey(uid), coldkey=get_mock_coldkey(uid), **kwargs
     )
+
+
+class FakeWebsocket(ClientConnection):
+    def __init__(self, *args, seed, **kwargs):
+        protocol = ClientProtocol(parse_uri("ws://127.0.0.1:9945"))
+        super().__init__(socket=None, protocol=protocol, **kwargs)
+        self.seed = seed
+        self.received = deque()
+
+    def send(self, payload: str, *args, **kwargs):
+        received = json.loads(payload)
+        id_ = received.pop("id")
+        self.received.append((received, id_))
+
+    def recv(self, *args, **kwargs):
+        item, _id = self.received.pop()
+        try:
+            response = WEBSOCKET_RESPONSES[self.seed][item["method"]][
+                json.dumps(item["params"])
+            ]
+            response["id"] = _id
+            return json.dumps(response)
+        except (KeyError, TypeError):
+            print("ERROR", self.seed, item["method"], item["params"])
+            raise
+
+    def close(self, *args, **kwargs):
+        pass
diff --git a/tests/integration_tests/test_subtensor_integration.py b/tests/integration_tests/test_subtensor_integration.py
index af391259ba..760673d0a1 100644
--- a/tests/integration_tests/test_subtensor_integration.py
+++ b/tests/integration_tests/test_subtensor_integration.py
@@ -1,46 +1,10 @@
-import json
-from collections import deque
-from websockets.sync.client import ClientConnection, ClientProtocol
-from websockets.uri import parse_uri
-
 import pytest
 from bittensor.utils.balance import Balance
 from bittensor.core.chain_data.axon_info import AxonInfo
 
 from bittensor import NeuronInfo
 from bittensor.core.subtensor import Subtensor
-from tests.helpers.integration_websocket_data import WEBSOCKET_RESPONSES
-
-
-class FakeWebsocket(ClientConnection):
-    def __init__(self, *args, seed, **kwargs):
-        protocol = ClientProtocol(parse_uri("ws://127.0.0.1:9945"))
-        super().__init__(socket=None, protocol=protocol, **kwargs)
-        self.seed = seed
-        self.received = deque()
-
-    def send(self, payload: str, *args, **kwargs):
-        received = json.loads(payload)
-        id_ = received.pop("id")
-        self.received.append((received, id_))
-
-    def recv(self):
-        item, _id = self.received.pop()
-        try:
-            response = WEBSOCKET_RESPONSES[self.seed][item["method"]][
-                json.dumps(item["params"])
-            ]
-            response["id"] = _id
-            return json.dumps(response)
-        except KeyError:
-            print("ERROR", self.seed, item["method"], item["params"])
-            raise
-        except TypeError:
-            print("TypeError", response)
-            raise
-
-    def close(self, *args, **kwargs):
-        pass
+from tests.helpers.helpers import FakeWebsocket
 
 
 @pytest.fixture