From e88e54ea2eb53d951e3bcfdbb8287d560bb4c269 Mon Sep 17 00:00:00 2001 From: ibraheem-opentensor Date: Fri, 13 Dec 2024 16:46:04 -0800 Subject: [PATCH] Adds overwrite and updates tests for btwallet3 --- bittensor_cli/cli.py | 30 +++++++++++++++--------- bittensor_cli/src/commands/wallets.py | 20 ++++++++++------ requirements.txt | 2 +- tests/e2e_tests/test_wallet_creations.py | 11 +++++---- 4 files changed, 39 insertions(+), 24 deletions(-) diff --git a/bittensor_cli/cli.py b/bittensor_cli/cli.py index 892d56fa..71d86f97 100755 --- a/bittensor_cli/cli.py +++ b/bittensor_cli/cli.py @@ -132,15 +132,10 @@ class Options: ss58_address = typer.Option( None, "--ss58", "--ss58-address", help="The SS58 address of the coldkey." ) - overwrite_coldkey = typer.Option( + overwrite = typer.Option( False, - help="Overwrite the old coldkey with the newly generated coldkey.", - prompt=True, - ) - overwrite_hotkey = typer.Option( - False, - help="Overwrite the old hotkey with the newly generated hotkey.", - prompt=True, + "--overwrite/--no-overwrite", + help="Overwrite the existing wallet file with the new one.", ) network = typer.Option( None, @@ -1725,6 +1720,7 @@ def wallet_regen_coldkey( json: Optional[str] = Options.json, json_password: Optional[str] = Options.json_password, use_password: Optional[bool] = Options.use_password, + overwrite: bool = Options.overwrite, quiet: bool = Options.quiet, verbose: bool = Options.verbose, ): @@ -1770,6 +1766,7 @@ def wallet_regen_coldkey( json, json_password, use_password, + overwrite, ) ) @@ -1780,6 +1777,7 @@ def wallet_regen_coldkey_pub( wallet_hotkey: Optional[str] = Options.wallet_hotkey, public_key_hex: Optional[str] = Options.public_hex_key, ss58_address: Optional[str] = Options.ss58_address, + overwrite: bool = Options.overwrite, quiet: bool = Options.quiet, verbose: bool = Options.verbose, ): @@ -1826,7 +1824,7 @@ def wallet_regen_coldkey_pub( rich.print("[red]Error: Invalid SS58 address or public key![/red]") raise typer.Exit() return self._run_command( - wallets.regen_coldkey_pub(wallet, ss58_address, public_key_hex) + wallets.regen_coldkey_pub(wallet, ss58_address, public_key_hex, overwrite) ) def wallet_regen_hotkey( @@ -1842,6 +1840,7 @@ def wallet_regen_hotkey( False, # Overriden to False help="Set to 'True' to protect the generated Bittensor key with a password.", ), + overwrite: bool = Options.overwrite, quiet: bool = Options.quiet, verbose: bool = Options.verbose, ): @@ -1880,6 +1879,7 @@ def wallet_regen_hotkey( json, json_password, use_password, + overwrite, ) ) @@ -1898,6 +1898,7 @@ def wallet_new_hotkey( False, # Overriden to False help="Set to 'True' to protect the generated Bittensor key with a password.", ), + overwrite: bool = Options.overwrite, quiet: bool = Options.quiet, verbose: bool = Options.verbose, ): @@ -1935,7 +1936,9 @@ def wallet_new_hotkey( validate=WV.WALLET, ) n_words = get_n_words(n_words) - return self._run_command(wallets.new_hotkey(wallet, n_words, use_password)) + return self._run_command( + wallets.new_hotkey(wallet, n_words, use_password, overwrite) + ) def wallet_new_coldkey( self, @@ -1949,6 +1952,7 @@ def wallet_new_coldkey( help="The number of words used in the mnemonic. Options: [12, 15, 18, 21, 24]", ), use_password: Optional[bool] = Options.use_password, + overwrite: bool = Options.overwrite, quiet: bool = Options.quiet, verbose: bool = Options.verbose, ): @@ -1985,7 +1989,9 @@ def wallet_new_coldkey( validate=WV.NONE, ) n_words = get_n_words(n_words) - return self._run_command(wallets.new_coldkey(wallet, n_words, use_password)) + return self._run_command( + wallets.new_coldkey(wallet, n_words, use_password, overwrite) + ) def wallet_check_ck_swap( self, @@ -2019,6 +2025,7 @@ def wallet_create_wallet( wallet_hotkey: Optional[str] = Options.wallet_hotkey, n_words: Optional[int] = None, use_password: bool = Options.use_password, + overwrite: bool = Options.overwrite, quiet: bool = Options.quiet, verbose: bool = Options.verbose, ): @@ -2064,6 +2071,7 @@ def wallet_create_wallet( wallet, n_words, use_password, + overwrite, ) ) diff --git a/bittensor_cli/src/commands/wallets.py b/bittensor_cli/src/commands/wallets.py index 9a8b0bfa..d49c03e2 100644 --- a/bittensor_cli/src/commands/wallets.py +++ b/bittensor_cli/src/commands/wallets.py @@ -69,6 +69,7 @@ async def regen_coldkey( json_path: Optional[str] = None, json_password: Optional[str] = "", use_password: Optional[bool] = True, + overwrite: Optional[bool] = False, ): """Creates a new coldkey under this wallet""" json_str: Optional[str] = None @@ -83,7 +84,7 @@ async def regen_coldkey( seed=seed, json=(json_str, json_password) if all([json_str, json_password]) else None, use_password=use_password, - overwrite=False, + overwrite=overwrite, ) if isinstance(new_wallet, Wallet): @@ -101,13 +102,14 @@ async def regen_coldkey_pub( wallet: Wallet, ss58_address: str, public_key_hex: str, + overwrite: Optional[bool] = False, ): """Creates a new coldkeypub under this wallet.""" try: new_coldkeypub = wallet.regenerate_coldkeypub( ss58_address=ss58_address, public_key=public_key_hex, - overwrite=False, + overwrite=overwrite, ) if isinstance(new_coldkeypub, Wallet): console.print( @@ -125,6 +127,7 @@ async def regen_hotkey( json_path: Optional[str], json_password: Optional[str] = "", use_password: Optional[bool] = False, + overwrite: Optional[bool] = False, ): """Creates a new hotkey under this wallet.""" json_str: Optional[str] = None @@ -141,7 +144,7 @@ async def regen_hotkey( seed=seed, json=(json_str, json_password) if all([json_str, json_password]) else None, use_password=use_password, - overwrite=False, + overwrite=overwrite, ) if isinstance(new_hotkey, Wallet): console.print( @@ -158,13 +161,14 @@ async def new_hotkey( wallet: Wallet, n_words: int, use_password: bool, + overwrite: Optional[bool] = False, ): """Creates a new hotkey under this wallet.""" try: wallet.create_new_hotkey( n_words=n_words, use_password=use_password, - overwrite=False, + overwrite=overwrite, ) except KeyFileError: print_error("KeyFileError: File is not writable") @@ -174,13 +178,14 @@ async def new_coldkey( wallet: Wallet, n_words: int, use_password: bool, + overwrite: Optional[bool] = False, ): """Creates a new coldkey under this wallet.""" try: wallet.create_new_coldkey( n_words=n_words, use_password=use_password, - overwrite=False, + overwrite=overwrite, ) except KeyFileError: print_error("KeyFileError: File is not writable") @@ -190,13 +195,14 @@ async def wallet_create( wallet: Wallet, n_words: int = 12, use_password: bool = True, + overwrite: Optional[bool] = False, ): """Creates a new wallet.""" try: wallet.create_new_coldkey( n_words=n_words, use_password=use_password, - overwrite=False, + overwrite=overwrite, ) except KeyFileError: print_error("KeyFileError: File is not writable") @@ -205,7 +211,7 @@ async def wallet_create( wallet.create_new_hotkey( n_words=n_words, use_password=False, - overwrite=False, + overwrite=overwrite, ) except KeyFileError: print_error("KeyFileError: File is not writable") diff --git a/requirements.txt b/requirements.txt index 4aa89c36..0a6f921c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,5 +16,5 @@ scalecodec==1.2.11 substrate-interface~=1.7.9 typer~=0.12 websockets>=14.1 -bittensor-wallet>=2.1.3 +bittensor-wallet>=3.0.0 bt-decode==0.4.0 \ No newline at end of file diff --git a/tests/e2e_tests/test_wallet_creations.py b/tests/e2e_tests/test_wallet_creations.py index 23b1e18f..7490a480 100644 --- a/tests/e2e_tests/test_wallet_creations.py +++ b/tests/e2e_tests/test_wallet_creations.py @@ -304,7 +304,7 @@ def test_wallet_creations(wallet_setup): assert wallet_status, message -def test_wallet_regen(wallet_setup): +def test_wallet_regen(wallet_setup, capfd): """ Test the regeneration of coldkeys, hotkeys, and coldkeypub files using mnemonics or ss58 address. @@ -342,7 +342,8 @@ def test_wallet_regen(wallet_setup): # Verify the command has output, as expected assert result.stdout is not None - mnemonics = extract_mnemonics_from_commands(result.stdout) + captured = capfd.readouterr() + mnemonics = extract_mnemonics_from_commands(captured.out) wallet_status, message = verify_wallet_dir( wallet_path, @@ -372,8 +373,8 @@ def test_wallet_regen(wallet_setup): "--mnemonic", mnemonics["coldkey"], "--no-use-password", + "--overwrite" ], - inputs=["y", "y"], ) # Wait a bit to ensure file system updates modification time @@ -412,8 +413,8 @@ def test_wallet_regen(wallet_setup): wallet_path, "--ss58-address", ss58_address, + "--overwrite" ], - inputs=["y"], ) # Wait a bit to ensure file system updates modification time @@ -447,8 +448,8 @@ def test_wallet_regen(wallet_setup): "--mnemonic", mnemonics["hotkey"], "--no-use-password", + "--overwrite", ], - inputs=["y"], ) # Wait a bit to ensure file system updates modification time