feat: Refactor phase 2 overview cmd & add test cov. Adds factories

bittensor/commands/overview.py

@@ -119,7 +119,7 @@ def _get_total_balance(
                 )
             if not coldkey_wallet.coldkeypub_file.exists_on_device():
                 console.print("[bold red]No wallets found.")
-                return
+                return [], None
             all_hotkeys = get_hotkey_wallets_for_wallet(coldkey_wallet)
 
         return all_hotkeys, total_balance
@@ -159,6 +159,26 @@ def _get_key_address(all_hotkeys: List["bittensor.wallet"]):
 
         return all_hotkey_addresses, hotkey_coldkey_to_hotkey_wallet
 
+    @staticmethod
+    def _process_neuron_results(
+        results: List[Tuple[int, List["bittensor.NeuronInfoLite"], Optional[str]]],
+        neurons: Dict[str, List["bittensor.NeuronInfoLite"]],
+        netuids: List[int],
+    ) -> Dict[str, List["bittensor.NeuronInfoLite"]]:
+        for result in results:
+            netuid, neurons_result, err_msg = result
+            if err_msg is not None:
+                console.print(f"netuid '{netuid}': {err_msg}")
+
+            if len(neurons_result) == 0:
+                # Remove netuid from overview if no neurons are found.
+                netuids.remove(netuid)
+                del neurons[str(netuid)]
+            else:
+                # Add neurons to overview.
+                neurons[str(netuid)] = neurons_result
+        return neurons
+
     def _run(cli: "bittensor.cli", subtensor: "bittensor.subtensor"):
         r"""Prints an overview for the wallet's colkey."""
         console = bittensor.__console__
@@ -226,18 +246,9 @@ def _run(cli: "bittensor.cli", subtensor: "bittensor.subtensor"):
                 )
                 executor.shutdown(wait=True)  # wait for all complete
 
-                for result in results:
-                    netuid, neurons_result, err_msg = result
-                    if err_msg is not None:
-                        console.print(f"netuid '{netuid}': {err_msg}")
-
-                    if len(neurons_result) == 0:
-                        # Remove netuid from overview if no neurons are found.
-                        netuids.remove(netuid)
-                        del neurons[str(netuid)]
-                    else:
-                        # Add neurons to overview.
-                        neurons[str(netuid)] = neurons_result
+                neurons = OverviewCommand._process_neuron_results(
+                    results, neurons, netuids
+                )
 
             total_coldkey_stake_from_metagraph = defaultdict(
                 lambda: bittensor.Balance(0.0)

requirements/dev.txt

@@ -12,3 +12,4 @@ mypy==1.8.0
 types-retry==0.9.9.4
 freezegun==1.5.0
 torch>=1.13.1
+factory-boy==3.3.0

tests/unit_tests/factories/neuron_factory.py

@@ -0,0 +1,63 @@
+import factory
+
+from bittensor.chain_data import AxonInfo, NeuronInfoLite, PrometheusInfo
+from bittensor.utils.balance import Balance
+
+
+class BalanceFactory(factory.Factory):
+    class Meta:
+        model = Balance
+
+    balance = factory.Faker("pyfloat", left_digits=3, right_digits=6, positive=True)
+
+
+class PrometheusInfoFactory(factory.Factory):
+    class Meta:
+        model = PrometheusInfo
+
+    block = factory.Faker("random_int", min=0, max=100)
+    version = factory.Faker("random_int", min=0, max=100)
+    ip = factory.Faker("ipv4")
+    port = factory.Faker("random_int", min=0, max=100)
+    ip_type = factory.Faker("random_int", min=0, max=100)
+
+
+class AxonInfoFactory(factory.Factory):
+    class Meta:
+        model = AxonInfo
+
+    version = factory.Faker("random_int", min=0, max=100)
+    ip = factory.Faker("ipv4")
+    port = factory.Faker("random_int", min=0, max=100)
+    ip_type = factory.Faker("random_int", min=0, max=100)
+    hotkey = factory.Faker("uuid4")
+    coldkey = factory.Faker("uuid4")
+
+
+class NeuronInfoLiteFactory(factory.Factory):
+    class Meta:
+        model = NeuronInfoLite
+
+    hotkey = factory.Faker("uuid4")
+    coldkey = factory.Faker("uuid4")
+    uid = factory.Sequence(lambda n: n)
+    netuid = factory.Sequence(lambda n: n)
+    active = factory.Faker("random_int", min=0, max=1)
+    stake = factory.SubFactory(BalanceFactory)
+    stake_dict = factory.Dict({"balance": 10})
+    total_stake = factory.SubFactory(BalanceFactory)
+    rank = factory.Faker("pyfloat", left_digits=3, right_digits=6, positive=True)
+    emission = factory.Faker("pyfloat", left_digits=3, right_digits=6, positive=True)
+    incentive = factory.Faker("pyfloat", left_digits=3, right_digits=6, positive=True)
+    consensus = factory.Faker("pyfloat", left_digits=3, right_digits=6, positive=True)
+    trust = factory.Faker("pyfloat", left_digits=3, right_digits=6, positive=True)
+    validator_trust = factory.Faker(
+        "pyfloat", left_digits=3, right_digits=6, positive=True
+    )
+    dividends = factory.Faker("pyfloat", left_digits=3, right_digits=6, positive=True)
+    last_update = factory.Faker("unix_time")
+    validator_permit = factory.Faker("boolean")
+    prometheus_info = factory.SubFactory(PrometheusInfoFactory)
+    axon_info = factory.SubFactory(AxonInfoFactory)
+    pruning_score = factory.Faker("random_int", min=0, max=100)
+    is_null = factory.Faker("boolean")

tests/unit_tests/test_overview.py

@@ -8,6 +8,7 @@
 # Bittensor
 import bittensor
 from bittensor.commands.overview import OverviewCommand
+from tests.unit_tests.factories.neuron_factory import NeuronInfoLiteFactory
 
 
 @pytest.fixture
@@ -183,3 +184,83 @@ def test_get_hotkeys_error():
     # Act
     with pytest.raises(TypeError):
         OverviewCommand._get_hotkeys(cli, all_hotkeys)
+
+
+@pytest.fixture
+def neuron_info():
+    return [
+        (1, [NeuronInfoLiteFactory(netuid=1)], None),
+        (2, [NeuronInfoLiteFactory(netuid=2)], None),
+    ]
+
+
+@pytest.fixture
+def neurons_dict():
+    return {
+        "1": [NeuronInfoLiteFactory(netuid=1)],
+        "2": [NeuronInfoLiteFactory(netuid=2)],
+    }
+
+
+@pytest.fixture
+def netuids_list():
+    return [1, 2]
+
+
+# Test cases
+@pytest.mark.parametrize(
+    "test_id, results, expected_neurons, expected_netuids",
+    [
+        # Test ID: 01 - Happy path, all neurons processed correctly
+        (
+            "01",
+            [
+                (1, [NeuronInfoLiteFactory(netuid=1)], None),
+                (2, [NeuronInfoLiteFactory(netuid=2)], None),
+            ],
+            {
+                "1": [NeuronInfoLiteFactory(netuid=1)],
+                "2": [NeuronInfoLiteFactory(netuid=2)],
+            },
+            [1, 2],
+        ),
+        # Test ID: 02 - Error message present, should skip processing for that netuid
+        (
+            "02",
+            [
+                (1, [NeuronInfoLiteFactory(netuid=1)], None),
+                (2, [], "Error fetching data"),
+            ],
+            {"1": [NeuronInfoLiteFactory()]},
+            [1],
+        ),
+        # Test ID: 03 - No neurons found for a netuid, should remove the netuid
+        (
+            "03",
+            [(1, [NeuronInfoLiteFactory()], None), (2, [], None)],
+            {"1": [NeuronInfoLiteFactory()]},
+            [1],
+        ),
+        # Test ID: 04 - Mixed conditions
+        (
+            "04",
+            [
+                (1, [NeuronInfoLiteFactory(netuid=1)], None),
+                (2, [], None),
+            ],
+            {"1": [NeuronInfoLiteFactory()]},
+            [1],
+        ),
+    ],
+)
+def test_process_neuron_results(
+    test_id, results, expected_neurons, expected_netuids, neurons_dict, netuids_list
+):
+    # Act
+    actual_neurons = OverviewCommand._process_neuron_results(
+        results, neurons_dict, netuids_list
+    )
+
+    # Assert
+    assert actual_neurons.keys() == expected_neurons.keys(), f"Failed test {test_id}"
+    assert netuids_list == expected_netuids, f"Failed test {test_id}"