diff --git a/bittensor_cli/src/bittensor/subtensor_interface.py b/bittensor_cli/src/bittensor/subtensor_interface.py index 65474b324..c72d3d809 100644 --- a/bittensor_cli/src/bittensor/subtensor_interface.py +++ b/bittensor_cli/src/bittensor/subtensor_interface.py @@ -1327,20 +1327,29 @@ async def get_stake_for_coldkeys( This function is useful for analyzing the stake distribution and delegation patterns of multiple accounts simultaneously, offering a broader perspective on network participation and investment strategies. """ - result = await self.query_runtime_api( - runtime_api="StakeInfoRuntimeApi", - method="get_stake_info_for_coldkeys", - params=[coldkey_ss58_list], - block_hash=block_hash, - ) - if result is None: - return None - + BATCH_SIZE = 60 + + tasks = [] + for i in range(0, len(coldkey_ss58_list), BATCH_SIZE): + ss58_chunk = coldkey_ss58_list[i : i + BATCH_SIZE] + tasks.append( + self.query_runtime_api( + runtime_api="StakeInfoRuntimeApi", + method="get_stake_info_for_coldkeys", + params=[ss58_chunk], + block_hash=block_hash, + ) + ) + results = await asyncio.gather(*tasks) stake_info_map = {} - for coldkey_bytes, stake_info_list in result: - coldkey_ss58 = decode_account_id(coldkey_bytes) - stake_info_map[coldkey_ss58] = StakeInfo.list_from_any(stake_info_list) - return stake_info_map + for result in results: + if result is None: + continue + for coldkey_bytes, stake_info_list in result: + coldkey_ss58 = decode_account_id(coldkey_bytes) + stake_info_map[coldkey_ss58] = StakeInfo.list_from_any(stake_info_list) + + return stake_info_map if stake_info_map else None async def all_subnets(self, block_hash: Optional[str] = None) -> list[DynamicInfo]: result = await self.query_runtime_api( diff --git a/tests/e2e_tests/test_wallet_creations.py b/tests/e2e_tests/test_wallet_creations.py index 52aa31399..ac82ec2a4 100644 --- a/tests/e2e_tests/test_wallet_creations.py +++ b/tests/e2e_tests/test_wallet_creations.py @@ -461,3 +461,60 @@ def test_wallet_regen(wallet_setup, capfd): initial_hotkey_mod_time != new_hotkey_mod_time ), "Hotkey file was not regenerated as expected" print("Passed wallet regen_hotkey command ✅") + + +def test_wallet_balance_all(local_chain, wallet_setup, capfd): + """ + Test the wallet balance --all command with a large number of wallets. + + Steps: + 1. Create 100 wallets + 2. Run wallet balance --all command + 3. Verify the output contains all wallet names and their balances + + Raises: + AssertionError: If any of the checks or verifications fail + """ + wallet_path_name = "//Alice" + keypair, wallet, wallet_path, exec_command = wallet_setup(wallet_path_name) + + print("Creating 100 wallets for testing balance --all command 🧪") + num_wallets = 100 + wallet_names = [] + + for i in range(num_wallets): + wallet_name = f"test_wallet_{i}" + wallet_names.append(wallet_name) + + result = exec_command( + command="wallet", + sub_command="new-coldkey", + extra_args=[ + "--wallet-name", + wallet_name, + "--wallet-path", + wallet_path, + "--n-words", + "12", + "--no-use-password", + ], + ) + + wallet_status, message = verify_wallet_dir(wallet_path, wallet_name) + assert wallet_status, message + + print("Testing wallet balance --all command 🧪") + result = exec_command( + command="wallet", + sub_command="balance", + extra_args=["--wallet-path", wallet_path, "--all"], + ) + + output = result.stdout + + for wallet_name in wallet_names: + assert ( + wallet_name in output + ), f"Wallet {wallet_name} not found in balance --all output" + + print("Passed wallet balance --all command with 100 wallets ✅")