add force_register_neuron method

Author: Roman

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]]: