Adds limit of ss58 addresses per call when fetching total_stake

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(

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 ✅")