wallet_setup fixture + fresh wallets

Author: ibraheem-opentensor

tests/e2e_tests/__init__.py

@@ -2,16 +2,18 @@
 import os
 import re
 import shlex
+import shutil
 import signal
 import subprocess
 import time
 
 import pytest
-
 from bittensor_cli.src.bittensor.async_substrate_interface import (
     AsyncSubstrateInterface,
 )
 
+from .utils import setup_wallet
+
 
 # Fixture for setting up and tearing down a localnet.sh chain between tests
 @pytest.fixture(scope="function")
@@ -70,3 +72,19 @@ def wait_for_node_start(process, pattern):
 
     # Ensure the process has terminated
     process.wait()
+
+
+@pytest.fixture(scope="function")
+def wallet_setup():
+    wallet_paths = []
+
+    def _setup_wallet(uri: str):
+        keypair, wallet, wallet_path, exec_command = setup_wallet(uri)
+        wallet_paths.append(wallet_path)
+        return keypair, wallet, wallet_path, exec_command
+
+    yield _setup_wallet
+
+    # Cleanup after the test
+    for path in wallet_paths:
+        shutil.rmtree(path, ignore_errors=True)

tests/e2e_tests/conftest.py

@@ -2,8 +2,6 @@
 
 from bittensor_cli.src.bittensor.balances import Balance
 
-from tests.e2e_tests.utils import setup_wallet, remove_wallets
-
 """
 Verify commands:
 
@@ -17,7 +15,7 @@
 """
 
 
-def test_root_commands(local_chain):
+def test_root_commands(local_chain, wallet_setup):
     """
     Test the root commands and inspects their output
 
@@ -38,10 +36,10 @@ def test_root_commands(local_chain):
     wallet_path_bob = "//Bob"
 
     # Create wallet for Alice
-    keypair_alice, wallet_alice, wallet_path_alice, exec_command_alice = setup_wallet(
+    keypair_alice, wallet_alice, wallet_path_alice, exec_command_alice = wallet_setup(
         wallet_path_alice
     )
-    keypair_bob, wallet_bob, wallet_path_bob, exec_command_bob = setup_wallet(
+    keypair_bob, wallet_bob, wallet_path_bob, exec_command_bob = wallet_setup(
         wallet_path_bob
     )
 
@@ -260,5 +258,3 @@ def test_root_commands(local_chain):
     assert "✅ Finalized" in undelegate_alice.stdout
 
     print("✅ Passed Root commands")
-    remove_wallets(wallet_path_alice)
-    remove_wallets(wallet_path_bob)

tests/e2e_tests/test_root.py

@@ -2,7 +2,6 @@
 import time
 
 from bittensor_cli.src.bittensor.balances import Balance
-from tests.e2e_tests.utils import setup_wallet, remove_wallets
 
 """
 Verify commands:
@@ -17,7 +16,7 @@
 """
 
 
-def test_staking(local_chain):
+def test_staking(local_chain, wallet_setup):
     """
     Test staking & sudo commands and inspect their output
 
@@ -37,7 +36,7 @@ def test_staking(local_chain):
     wallet_path_alice = "//Alice"
 
     # Create wallet for Alice
-    keypair_alice, wallet_alice, wallet_path_alice, exec_command_alice = setup_wallet(
+    keypair_alice, wallet_alice, wallet_path_alice, exec_command_alice = wallet_setup(
         wallet_path_alice
     )
 
@@ -224,4 +223,3 @@ def test_staking(local_chain):
         10
     )
     print("✅ Passed staking and sudo commands")
-    remove_wallets(wallet_path_alice)  

tests/e2e_tests/test_staking_sudo.py

@@ -3,8 +3,6 @@
 import time
 from typing import Dict, Optional, Tuple
 
-from tests.e2e_tests.utils import setup_wallet, remove_wallets, get_path_from_uri
-
 """
 Verify commands:
 
@@ -158,7 +156,7 @@ def extract_mnemonics_from_commands(output: str) -> Dict[str, Optional[str]]:
     return mnemonics
 
 
-def test_wallet_creations():
+def test_wallet_creations(wallet_setup):
     """
     Test the creation and verification of wallet keys and directories in the Bittensor network.
 
@@ -175,7 +173,7 @@ def test_wallet_creations():
     """
 
     wallet_path_name = "//Alice"
-    keypair, wallet, wallet_path, exec_command = setup_wallet(wallet_path_name)
+    keypair, wallet, wallet_path, exec_command = wallet_setup(wallet_path_name)
 
     result = exec_command(
         command="wallet", sub_command="list", extra_args=["--wallet-path", wallet_path]
@@ -296,10 +294,9 @@ def test_wallet_creations():
         wallet_path, "new_coldkey", hotkey_name="new_hotkey"
     )
     assert wallet_status, message
-    remove_wallets(wallet_path)
 
 
-def test_wallet_regen():
+def test_wallet_regen(wallet_setup):
     """
     Test the regeneration of coldkeys, hotkeys, and coldkeypub files using mnemonics or ss58 address.
 
@@ -313,7 +310,7 @@ def test_wallet_regen():
         AssertionError: If any of the checks or verifications fail
     """
     wallet_path_name = "//Bob"
-    keypair, wallet, wallet_path, exec_command = setup_wallet(wallet_path_name)
+    keypair, wallet, wallet_path, exec_command = wallet_setup(wallet_path_name)
 
     # Create a new wallet (coldkey + hotkey)
     result = exec_command(
@@ -450,4 +447,3 @@ def test_wallet_regen():
         initial_hotkey_mod_time != new_hotkey_mod_time
     ), "Hotkey file was not regenerated as expected"
     print("Passed wallet regen_hotkey command ✅")
-    remove_wallets(wallet_path)

tests/e2e_tests/test_wallet_creations.py

@@ -1,14 +1,9 @@
-import logging
-
 from bittensor_cli.src.bittensor.balances import Balance
-
 from tests.e2e_tests.utils import (
     extract_coldkey_balance,
-    setup_wallet,
     validate_wallet_inspect,
     validate_wallet_overview,
     verify_subnet_entry,
-    remove_wallets,
 )
 
 """
@@ -22,7 +17,7 @@
 """
 
 
-def test_wallet_overview_inspect(local_chain):
+def test_wallet_overview_inspect(local_chain, wallet_setup):
     """
     Test the overview and inspect commands of the wallet by interaction with subnets
 
@@ -40,7 +35,7 @@ def test_wallet_overview_inspect(local_chain):
     wallet_path_name = "//Alice"
 
     # Create wallet for Alice
-    keypair, wallet, wallet_path, exec_command = setup_wallet(wallet_path_name)
+    keypair, wallet, wallet_path, exec_command = wallet_setup(wallet_path_name)
 
     # Register a subnet with sudo as Alice
     result = exec_command(
@@ -176,7 +171,6 @@ def test_wallet_overview_inspect(local_chain):
         ],  # (netuid, hotkey-display, stake, check_emissions)
     )
     print("Passed wallet overview, inspect command ✅")
-    remove_wallets(wallet_path)
 
 
 """
@@ -187,7 +181,7 @@ def test_wallet_overview_inspect(local_chain):
 """
 
 
-def test_wallet_transfer(local_chain):
+def test_wallet_transfer(local_chain, wallet_setup):
     """
     Test the transfer and balance functionality in the Bittensor network.
 
@@ -206,10 +200,10 @@ def test_wallet_transfer(local_chain):
     wallet_path_bob = "//Bob"
 
     # Create wallets for Alice and Bob
-    keypair_alice, wallet_alice, wallet_path_alice, exec_command_alice = setup_wallet(
+    keypair_alice, wallet_alice, wallet_path_alice, exec_command_alice = wallet_setup(
         wallet_path_alice
     )
-    keypair_bob, wallet_bob, wallet_path_bob, exec_command_bob = setup_wallet(
+    keypair_bob, wallet_bob, wallet_path_bob, exec_command_bob = wallet_setup(
         wallet_path_bob
     )
 
@@ -332,7 +326,7 @@ def test_wallet_transfer(local_chain):
 
     wallet_path_anakin = "//Anakin"
     keypair_anakin, wallet_anakin, wallet_path_anakin, exec_command_anakin = (
-        setup_wallet(wallet_path_anakin)
+        wallet_setup(wallet_path_anakin)
     )
 
     # Attempt transferring to Alice
@@ -358,8 +352,4 @@ def test_wallet_transfer(local_chain):
 
     # This transfer is expected to fail due to low balance
     assert "❌ Not enough balance" in result.stdout
-    print("Testing wallet transfer, balance command ✅")
-
-    remove_wallets(wallet_path_alice)
-    remove_wallets(wallet_path_bob)
-    remove_wallets(wallet_path_anakin)
+    print("✅Passed wallet transfer, balance command ")

tests/e2e_tests/test_wallet_interactions.py

@@ -5,23 +5,18 @@
 import sys
 from typing import List, Tuple
 
+from bittensor_cli.cli import CLIManager
 from bittensor_wallet import Wallet
 from substrateinterface import Keypair
 from typer.testing import CliRunner
 
-from bittensor_cli.cli import CLIManager
-
 template_path = os.getcwd() + "/neurons/"
 templates_repo = "templates repository"
 
 
-def get_path_from_uri(uri: str):
-    return f"/tmp/btcli-e2e-wallet-{uri.strip('/')}"
-
-
 def setup_wallet(uri: str):
     keypair = Keypair.create_from_uri(uri)
-    wallet_path = get_path_from_uri(uri)
+    wallet_path = f"/tmp/btcli-e2e-wallet-{uri.strip('/')}"
     wallet = Wallet(path=wallet_path)
     wallet.set_coldkey(keypair=keypair, encrypt=False, overwrite=True)
     wallet.set_coldkeypub(keypair=keypair, encrypt=False, overwrite=True)
@@ -287,7 +282,3 @@ def uninstall_templates(install_dir):
     )
     # Delete everything in directory
     shutil.rmtree(install_dir)
-
-
-def remove_wallets(wallets_dir):
-    shutil.rmtree(wallets_dir)