diff --git a/bittensor/keyfile.py b/bittensor/keyfile.py index f0962c09a4..b5157cea4a 100644 --- a/bittensor/keyfile.py +++ b/bittensor/keyfile.py @@ -276,13 +276,12 @@ def get_coldkey_password_from_environment(coldkey_name: str) -> Optional[str]: Returns: password (str): The password retrieved from the environment variables, or ``None`` if not found. """ - for env_var in os.environ: - if env_var.upper().startswith("BT_COLD_PW_") and env_var.upper().endswith( - coldkey_name.upper() - ): - return os.getenv(env_var) - - return None + envs = { + normalized_env_name: env_value + for env_name, env_value in os.environ.items() + if (normalized_env_name := env_name.upper()).startswith("BT_COLD_PW_") + } + return envs.get(f"BT_COLD_PW_{coldkey_name.upper()}") def decrypt_keyfile_data( diff --git a/tests/unit_tests/test_keyfile.py b/tests/unit_tests/test_keyfile.py index 266b234062..d20af809f9 100644 --- a/tests/unit_tests/test_keyfile.py +++ b/tests/unit_tests/test_keyfile.py @@ -28,6 +28,8 @@ from substrateinterface.exceptions import ConfigurationError from bip39 import bip39_validate +from bittensor import get_coldkey_password_from_environment + def test_generate_mnemonic(): """ @@ -606,3 +608,18 @@ def test_deserialize_keypair_from_keyfile_data(keyfile_setup_teardown): assert deserialized_keypair.ss58_address == keypair.ss58_address assert deserialized_keypair.public_key == keypair.public_key assert deserialized_keypair.private_key == keypair.private_key + + +def test_get_coldkey_password_from_environment(monkeypatch): + password_by_wallet = { + "WALLET": "password", + "my_wallet": "password", + } + + monkeypatch.setenv("bt_cold_pw_wallet", password_by_wallet["WALLET"]) + monkeypatch.setenv("BT_COLD_PW_My_Wallet", password_by_wallet["my_wallet"]) + + for wallet, password in password_by_wallet.items(): + assert get_coldkey_password_from_environment(wallet) == password + + assert get_coldkey_password_from_environment("non_existent_wallet") is None