Fix tests to be ready for rust-based bittensor-wallet (#2336)

* replace `from substrateinterface import Keypair` to `from bittensor_wallet import Keypair`

* apply raise to the test

* use Keypair from bittensor_wallet for whole repo

* change test name

* ruff linting

* set wallet to v 1.0.0

* Bittensor wallet -> 2.0.0

* test Axon online debug

* test Axon online debug 2

* test Axon online debug 3

* test Axon online debug 4

* test Axon online debug 5

* test Axon online debug 6

* fixed tests, ruff

---------

Co-authored-by: opendansor <daniel@opentensor.dev>
Co-authored-by: ibraheem-opentensor <ibraheem@opentensor.dev>
Co-authored-by: ibraheem-opentensor <165814940+ibraheem-opentensor@users.noreply.github.com>

bittensor/core/axon.py

@@ -32,14 +32,15 @@
 from typing import Any, Awaitable, Callable, Optional, Tuple
 
 import uvicorn
-from bittensor_wallet import Wallet
+from bittensor_wallet import Wallet, Keypair
+
 from fastapi import APIRouter, Depends, FastAPI
 from fastapi.responses import JSONResponse
 from fastapi.routing import serialize_response
 from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
 from starlette.requests import Request
 from starlette.responses import Response
-from substrateinterface import Keypair
+
 
 from bittensor.core.chain_data import AxonInfo
 from bittensor.core.config import Config

bittensor/core/dendrite.py

@@ -23,8 +23,7 @@
 from typing import Any, AsyncGenerator, Optional, Union, Type
 
 import aiohttp
-from bittensor_wallet import Wallet
-from substrateinterface import Keypair
+from bittensor_wallet import Keypair, Wallet
 
 from bittensor.core.axon import Axon
 from bittensor.core.chain_data import AxonInfo

bittensor/utils/__init__.py

@@ -16,10 +16,10 @@
 # DEALINGS IN THE SOFTWARE.
 
 import hashlib
-from typing import List, Dict, Literal, Union, Optional, TYPE_CHECKING
+from typing import Literal, Union, Optional, TYPE_CHECKING
 
 import scalecodec
-from substrateinterface import Keypair
+from bittensor_wallet import Keypair
 from substrateinterface.utils import ss58
 
 from bittensor.core.settings import SS58_FORMAT
@@ -245,7 +245,7 @@ def _is_valid_ed25519_pubkey(public_key: Union[str, bytes]) -> bool:
         else:
             raise ValueError("public_key must be a string or bytes")
 
-        keypair = Keypair(public_key=public_key, ss58_format=SS58_FORMAT)
+        keypair = Keypair(public_key=public_key)
 
         ss58_addr = keypair.ss58_address
         return ss58_addr is not None

bittensor/utils/deprecated.py

@@ -42,7 +42,7 @@
     Keyfile,
 )
 from bittensor_wallet.wallet import display_mnemonic_msg, Wallet  # noqa: F401
-from substrateinterface import Keypair  # noqa: F401
+from bittensor_wallet import Keypair  # noqa: F401
 
 from bittensor.core import settings
 from bittensor.core.axon import Axon

bittensor/utils/weight_utils.py

@@ -26,7 +26,7 @@
 
 from numpy.typing import NDArray
 from scalecodec import U16, ScaleBytes, Vec
-from substrateinterface import Keypair
+from bittensor_wallet import Keypair
 
 from bittensor.utils.btlogging import logging
 from bittensor.utils.registration import legacy_torch_api_compat, torch, use_torch

requirements/prod.txt

@@ -20,4 +20,4 @@ python-Levenshtein
 scalecodec==1.2.11
 substrate-interface~=1.7.9
 uvicorn
-bittensor-wallet==1.0.0
+bittensor-wallet==2.0.0

tests/e2e_tests/conftest.py

@@ -45,7 +45,10 @@ def local_chain(request):
 
     # install neuron templates
     logging.info("downloading and installing neuron templates from github")
-    templates_dir = clone_or_update_templates()
+    # commit with subnet-template-repo changes for rust wallet
+    templates_dir = clone_or_update_templates(
+        "334d3da101279218b3a4c9d72a12d517f6e39be3"
+    )
     install_templates(templates_dir)
 
     timestamp = int(time.time())

tests/e2e_tests/utils/e2e_test_utils.py

@@ -3,7 +3,7 @@
 import subprocess
 import sys
 
-from substrateinterface import Keypair
+from bittensor_wallet import Keypair
 
 import bittensor
 

tests/integration_tests/test_subtensor_integration.py

@@ -18,10 +18,12 @@
 import unittest
 from unittest.mock import MagicMock, patch
 
-from substrateinterface import Keypair
+import pytest
+from bittensor_wallet import Keypair
 
 import bittensor
 from bittensor.core import settings
+from bittensor.core.extrinsics import transfer
 from bittensor.utils.balance import Balance
 from bittensor.utils.mock import MockSubtensor
 from tests.helpers import (
@@ -30,7 +32,6 @@
     get_mock_keypair,
     get_mock_wallet,
 )
-from bittensor.core.extrinsics import transfer
 
 
 class TestSubtensor(unittest.TestCase):
@@ -171,7 +172,7 @@ def test_transfer_invalid_dest(self):
         )
         self.assertFalse(fail, msg="Transfer should fail because of invalid dest")
 
-    def test_transfer_dest_as_bytes(self):
+    def test_transfer_dest_as_bytes_fails(self):
         fake_coldkey = get_mock_coldkey(1)
         with patch(
             "bittensor.core.extrinsics.transfer.do_transfer",
@@ -183,13 +184,14 @@ def test_transfer_dest_as_bytes(self):
             self.subtensor.get_balance = MagicMock(return_value=self.balance)
 
             dest_as_bytes: bytes = Keypair(fake_coldkey).public_key
-            success = self.subtensor.transfer(
-                self.wallet,
-                dest_as_bytes,  # invalid dest
-                amount=200,
-                wait_for_inclusion=True,
-            )
-            self.assertTrue(success, msg="Transfer should succeed")
+
+            with pytest.raises(TypeError):
+                self.subtensor.transfer(
+                    self.wallet,
+                    dest_as_bytes,  # invalid dest
+                    amount=200,
+                    wait_for_inclusion=True,
+                )
 
     def test_set_weights(self):
         chain_weights = [0]

tests/unit_tests/test_axon.py

@@ -380,10 +380,16 @@ def test_to_string(info_return, expected_output, test_id):
 )
 def test_valid_ipv4_and_ipv6_address(ip, port, expected_ip_type, test_id):
     # Arrange
+    hotkey = MockHotkey("5EemgxS7cmYbD34esCFoBgUZZC8JdnGtQvV5Qw3QFUCRRtGP")
+    coldkey = MockHotkey("5EemgxS7cmYbD34esCFoBgUZZC8JdnGtQvV5Qw3QFUCRRtGP")
+    coldkeypub = MockHotkey("5EemgxS7cmYbD34esCFoBgUZZC8JdnGtQvV5Qw3QFUCRRtGP")
+    wallet = MockWallet(hotkey, coldkey, coldkeypub)
+
     axon = Axon()
     axon.ip = ip
     axon.external_ip = ip
     axon.port = port
+    axon.wallet = wallet
 
     # Act
     ip_type = axon.info().ip_type
@@ -586,7 +592,7 @@ async def test_ping__no_dendrite(self, http_client):
         assert (response.status_code, response.json()) == (
             401,
             {
-                "message": "Not Verified with error: No SS58 formatted address or public key provided"
+                "message": "Not Verified with error: No SS58 formatted address or public key provided."
             },
         )
 

tests/unit_tests/test_dendrite.py

@@ -73,7 +73,8 @@ def axon_info():
 
 @pytest.fixture(scope="session")
 def setup_axon():
-    axon = Axon()
+    wallet = get_mock_wallet()
+    axon = Axon(wallet)
     axon.attach(forward_fn=dummy)
     axon.start()
     yield axon