Add force_register_neuron into MockSubtensor

bittensor/utils/mock/subtensor_mock.py

@@ -1,6 +1,7 @@
 from collections.abc import Mapping
 from dataclasses import dataclass
 from hashlib import sha256
+from random import randint
 from types import SimpleNamespace
 from typing import Any, Optional, Union, TypedDict
 from unittest.mock import MagicMock, patch
@@ -358,6 +359,104 @@ def set_difficulty(self, netuid: int, difficulty: int) -> None:
 
         subtensor_state["Difficulty"][netuid][self.block_number] = difficulty
 
+    def _register_neuron(self, netuid: int, hotkey: str, coldkey: str) -> int:
+        subtensor_state = self.chain_state["SubtensorModule"]
+        if netuid not in subtensor_state["NetworksAdded"]:
+            raise Exception("Subnet does not exist")
+
+        subnetwork_n = self._get_most_recent_storage(
+            subtensor_state["SubnetworkN"][netuid]
+        )
+
+        if subnetwork_n > 0 and any(
+            self._get_most_recent_storage(subtensor_state["Keys"][netuid][uid])
+            == hotkey
+            for uid in range(subnetwork_n)
+        ):
+            # already_registered
+            raise Exception("Hotkey already registered")
+        else:
+            # Not found
+            if subnetwork_n >= self._get_most_recent_storage(
+                subtensor_state["MaxAllowedUids"][netuid]
+            ):
+                # Subnet full, replace neuron randomly
+                uid = randint(0, subnetwork_n - 1)
+            else:
+                # Subnet not full, add new neuron
+                # Append as next uid and increment subnetwork_n
+                uid = subnetwork_n
+                subtensor_state["SubnetworkN"][netuid][self.block_number] = (
+                    subnetwork_n + 1
+                )
+
+            subtensor_state["Stake"][hotkey] = {}
+            subtensor_state["Stake"][hotkey][coldkey] = {}
+            subtensor_state["Stake"][hotkey][coldkey][self.block_number] = 0
+
+            subtensor_state["Uids"][netuid][hotkey] = {}
+            subtensor_state["Uids"][netuid][hotkey][self.block_number] = uid
+
+            subtensor_state["Keys"][netuid][uid] = {}
+            subtensor_state["Keys"][netuid][uid][self.block_number] = hotkey
+
+            subtensor_state["Owner"][hotkey] = {}
+            subtensor_state["Owner"][hotkey][self.block_number] = coldkey
+
+            subtensor_state["Active"][netuid][uid] = {}
+            subtensor_state["Active"][netuid][uid][self.block_number] = True
+
+            subtensor_state["LastUpdate"][netuid][uid] = {}
+            subtensor_state["LastUpdate"][netuid][uid][self.block_number] = (
+                self.block_number
+            )
+
+            subtensor_state["Rank"][netuid][uid] = {}
+            subtensor_state["Rank"][netuid][uid][self.block_number] = 0.0
+
+            subtensor_state["Emission"][netuid][uid] = {}
+            subtensor_state["Emission"][netuid][uid][self.block_number] = 0.0
+
+            subtensor_state["Incentive"][netuid][uid] = {}
+            subtensor_state["Incentive"][netuid][uid][self.block_number] = 0.0
+
+            subtensor_state["Consensus"][netuid][uid] = {}
+            subtensor_state["Consensus"][netuid][uid][self.block_number] = 0.0
+
+            subtensor_state["Trust"][netuid][uid] = {}
+            subtensor_state["Trust"][netuid][uid][self.block_number] = 0.0
+
+            subtensor_state["ValidatorTrust"][netuid][uid] = {}
+            subtensor_state["ValidatorTrust"][netuid][uid][self.block_number] = 0.0
+
+            subtensor_state["Dividends"][netuid][uid] = {}
+            subtensor_state["Dividends"][netuid][uid][self.block_number] = 0.0
+
+            subtensor_state["PruningScores"][netuid][uid] = {}
+            subtensor_state["PruningScores"][netuid][uid][self.block_number] = 0.0
+
+            subtensor_state["ValidatorPermit"][netuid][uid] = {}
+            subtensor_state["ValidatorPermit"][netuid][uid][self.block_number] = False
+
+            subtensor_state["Weights"][netuid][uid] = {}
+            subtensor_state["Weights"][netuid][uid][self.block_number] = []
+
+            subtensor_state["Bonds"][netuid][uid] = {}
+            subtensor_state["Bonds"][netuid][uid][self.block_number] = []
+
+            subtensor_state["Axons"][netuid][hotkey] = {}
+            subtensor_state["Axons"][netuid][hotkey][self.block_number] = {}
+
+            subtensor_state["Prometheus"][netuid][hotkey] = {}
+            subtensor_state["Prometheus"][netuid][hotkey][self.block_number] = {}
+
+            if hotkey not in subtensor_state["IsNetworkMember"]:
+                subtensor_state["IsNetworkMember"][hotkey] = {}
+            subtensor_state["IsNetworkMember"][hotkey][netuid] = {}
+            subtensor_state["IsNetworkMember"][hotkey][netuid][self.block_number] = True
+
+            return uid
+
     @staticmethod
     def _convert_to_balance(balance: Union["Balance", float, int]) -> "Balance":
         if isinstance(balance, float):
@@ -368,6 +467,37 @@ def _convert_to_balance(balance: Union["Balance", float, int]) -> "Balance":
 
         return balance
 
+    def force_register_neuron(
+        self,
+        netuid: int,
+        hotkey: str,
+        coldkey: str,
+        stake: Union["Balance", float, int] = Balance(0),
+        balance: Union["Balance", float, int] = Balance(0),
+    ) -> int:
+        """
+        Force register a neuron on the mock chain, returning the UID.
+        """
+        stake = self._convert_to_balance(stake)
+        balance = self._convert_to_balance(balance)
+
+        subtensor_state = self.chain_state["SubtensorModule"]
+        if netuid not in subtensor_state["NetworksAdded"]:
+            raise Exception("Subnet does not exist")
+
+        uid = self._register_neuron(netuid=netuid, hotkey=hotkey, coldkey=coldkey)
+
+        subtensor_state["TotalStake"][self.block_number] = (
+            self._get_most_recent_storage(subtensor_state["TotalStake"]) + stake.rao
+        )
+        subtensor_state["Stake"][hotkey][coldkey][self.block_number] = stake.rao
+
+        if balance.rao > 0:
+            self.force_set_balance(coldkey, balance)
+        self.force_set_balance(coldkey, balance)
+
+        return uid
+
     def force_set_balance(
         self, ss58_address: str, balance: Union["Balance", float, int] = Balance(0)
     ) -> tuple[bool, Optional[str]]:

tests/integration_tests/test_metagraph_integration.py

@@ -1,10 +1,11 @@
+import os
 from unittest import mock
 
-import bittensor
 import torch
-import os
-from bittensor.utils.mock import MockSubtensor
+
+import bittensor
 from bittensor.core.metagraph import METAGRAPH_STATE_DICT_NDARRAY_KEYS, get_save_dir
+from bittensor.utils.mock import MockSubtensor
 
 _subtensor_mock: MockSubtensor = MockSubtensor()
 

tests/integration_tests/test_subtensor_integration.py

@@ -1,10 +1,10 @@
 import pytest
 
-from bittensor import NeuronInfo
-from bittensor.core.chain_data.axon_info import AxonInfo
+from bittensor.core.chain_data import AxonInfo, NeuronInfo
 from bittensor.core.subtensor import Subtensor
 from bittensor.utils.balance import Balance
 from tests.helpers.helpers import FakeWebsocket
+from bittensor.utils.mock.subtensor_mock import MockSubtensor
 
 
 @pytest.fixture
@@ -126,3 +126,27 @@ async def test_get_neuron_for_pubkey_and_subnet(mocker):
     assert isinstance(result.total_stake, Balance)
     assert isinstance(result.axon_info, AxonInfo)
     assert result.is_null is False
+
+
+def test_mock_subtensor_force_register_neuron():
+    """Tests the force_register_neuron method of the MockSubtensor class."""
+    # Preps
+    test_netuid = 1
+    subtensor = MockSubtensor()
+    subtensor.create_subnet(netuid=test_netuid)
+
+    uid1 = subtensor.force_register_neuron(test_netuid, "hk1", "cc1")
+    uid2 = subtensor.force_register_neuron(test_netuid, "hk2", "cc2")
+
+    # Calls
+    neurons = subtensor.neurons(test_netuid)
+    neuron1 = subtensor.neuron_for_uid(uid1, test_netuid)
+    neuron2 = subtensor.neuron_for_uid(uid2, test_netuid)
+
+    # Assertions
+    assert len(neurons) == 2
+    assert [neuron1, neuron2] == neurons
+    assert neuron1.hotkey == "hk1"
+    assert neuron1.coldkey == "cc1"
+    assert neuron2.hotkey == "hk2"
+    assert neuron2.coldkey == "cc2"