test: bring back old (sync subtensor) tests and fix them

bittensor/core/extrinsics/registration.py

@@ -11,12 +11,12 @@
 
 from bittensor.utils import unlock_key
 from bittensor.utils.btlogging import logging
-from bittensor.utils.registration import log_no_torch_error, torch
+from bittensor.utils.registration import create_pow, log_no_torch_error, torch
 
 if TYPE_CHECKING:
     from bittensor_wallet import Wallet
     from bittensor.core.subtensor import Subtensor
-    from bittensor.utils.registration.pow import POWSolution, create_pow
+    from bittensor.utils.registration.pow import POWSolution
 
 
 def _do_burned_register(

bittensor/core/extrinsics/staking.py

@@ -93,8 +93,10 @@ def _check_threshold_amount(
     if amount is None:
         # Stake it all.
         staking_balance = Balance.from_tao(old_balance.tao)
+    elif not isinstance(amount, Balance):
+        staking_balance = Balance.from_tao(amount)
     else:
-        staking_balance = Balance.from_tao(amount.tao)
+        staking_balance = amount
 
     # Leave existential balance to keep key alive.
     if staking_balance > old_balance - existential_deposit:

tests/helpers/helpers.py

@@ -1,18 +1,18 @@
 import asyncio
-from collections import deque
 import json
+import time
+from collections import deque
 from typing import Union
 
-from websockets.asyncio.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
 from bittensor_wallet.mock.wallet_mock import get_mock_wallet
+from websockets.asyncio.client import ClientConnection, ClientProtocol
+from websockets.uri import parse_uri
 
-from bittensor.utils.balance import Balance
 from bittensor.core.chain_data import AxonInfo, NeuronInfo, PrometheusInfo
+from bittensor.utils.balance import Balance
 from tests.helpers.integration_websocket_data import WEBSOCKET_RESPONSES, METADATA
 
 
@@ -118,17 +118,15 @@ def __init__(self, *args, seed, **kwargs):
         self.received = deque()
         self._lock = asyncio.Lock()
 
-    async def send(self, payload: str, *args, **kwargs):
+    def send(self, payload: str, *args, **kwargs):
         received = json.loads(payload)
         id_ = received.pop("id")
-        async with self._lock:
-            self.received.append((received, id_))
+        self.received.append((received, id_))
 
-    async def recv(self, *args, **kwargs):
+    def recv(self, *args, **kwargs):
         while len(self.received) == 0:
-            await asyncio.sleep(0.1)
-        async with self._lock:
-            item, _id = self.received.pop()
+            time.sleep(0.1)
+        item, _id = self.received.pop()
         try:
             if item["method"] == "state_getMetadata":
                 response = {"jsonrpc": "2.0", "id": _id, "result": METADATA}
@@ -142,5 +140,16 @@ async def recv(self, *args, **kwargs):
             print("ERROR", self.seed, item["method"], item["params"])
             raise
 
-    async def close(self, *args, **kwargs):
+    def close(self, *args, **kwargs):
+        pass
+
+
+class FakeConnectContextManager:
+    def __init__(self, seed):
+        self.seed = seed
+
+    def __enter__(self):
+        return FakeWebsocket(seed=self.seed)
+
+    def __exit__(self, exc_type, exc, tb):
         pass

tests/helpers/integration_websocket_data.py

@@ -6423,6 +6423,12 @@
             }
         },
         "system_chain": {"[]": {"jsonrpc": "2.0", "result": "Bittensor"}},
+        "chain_getBlockHash": {
+            "[3264143]": {
+                "jsonrpc": "2.0",
+                "result": None,
+            }
+        },
     },
     "min_allowed_weights": {
         "chain_getHead": {

tests/integration_tests/test_metagraph_integration.py

@@ -33,18 +33,14 @@ def test_sync_block_0(self):
         self.metagraph.sync(lite=True, block=0, subtensor=self.sub)
 
     def test_load_sync_save(self):
-        with mock.patch.object(
-            self.sub.async_subtensor, "neurons_lite", return_value=[]
-        ):
+        with mock.patch.object(self.sub, "neurons_lite", return_value=[]):
             self.metagraph.sync(lite=True, subtensor=self.sub)
             self.metagraph.save()
             self.metagraph.load()
             self.metagraph.save()
 
     def test_load_sync_save_from_torch(self):
-        with mock.patch.object(
-            self.sub.async_subtensor, "neurons_lite", return_value=[]
-        ):
+        with mock.patch.object(self.sub, "neurons_lite", return_value=[]):
             self.metagraph.sync(lite=True, subtensor=self.sub)
 
         def deprecated_save_torch(metagraph):

tests/integration_tests/test_subtensor_integration.py

@@ -1,14 +1,13 @@
-import asyncio
 import os.path
 
 import pytest
-from bittensor.utils.balance import Balance
-from bittensor.core.chain_data.axon_info import AxonInfo
+from bt_decode import PortableRegistry, MetadataV15
 
 from bittensor import NeuronInfo
+from bittensor.core.chain_data.axon_info import AxonInfo
 from bittensor.core.subtensor import Subtensor
-from bt_decode import PortableRegistry, MetadataV15
-from tests.helpers.helpers import FakeWebsocket
+from bittensor.utils.balance import Balance
+from tests.helpers.helpers import FakeConnectContextManager
 
 
 @pytest.fixture
@@ -32,12 +31,11 @@ async def prepare_test(mocker, seed):
             MetadataV15.decode_from_metadata_option(f.read())
         )
     subtensor = Subtensor("unknown", _mock=True)
-    mocker.patch.object(subtensor.substrate.ws, "ws", FakeWebsocket(seed=seed))
-    mocker.patch.object(subtensor.substrate.ws, "_initialized", True)
-    mocker.patch.object(subtensor.substrate._async_instance, "registry", registry)
-    subtensor.substrate.ws._receiving_task = asyncio.create_task(
-        subtensor.substrate.ws._start_receiving()
+    mocker.patch(
+        "async_substrate_interface.sync_substrate.connect",
+        mocker.Mock(return_value=FakeConnectContextManager(seed=seed)),
     )
+    mocker.patch.object(subtensor.substrate, "registry", registry)
     return subtensor