@@ -14,6 +14,7 @@ use sp_runtime::{
1414 BoundedVec , Percent ,
1515 traits:: { BlakeTwo256 , Hash } ,
1616} ;
17+ use sp_std:: collections:: btree_set:: BTreeSet ;
1718use sp_std:: vec;
1819use subtensor_runtime_common:: { AlphaCurrency , NetUid , TaoCurrency } ;
1920
@@ -1565,4 +1566,102 @@ mod pallet_benchmarks {
15651566 #[ extrinsic_call]
15661567 _ ( RawOrigin :: Signed ( coldkey. clone ( ) ) , netuid, hotkey. clone ( ) ) ;
15671568 }
1569+ #[ benchmark]
1570+ fn set_root_claim_type ( ) {
1571+ let coldkey: T :: AccountId = whitelisted_caller ( ) ;
1572+
1573+ #[ extrinsic_call]
1574+ _ ( RawOrigin :: Signed ( coldkey. clone ( ) ) , RootClaimTypeEnum :: Keep ) ;
1575+ }
1576+
1577+ #[ benchmark]
1578+ fn claim_root ( ) {
1579+ let coldkey: T :: AccountId = whitelisted_caller ( ) ;
1580+ let hotkey: T :: AccountId = account ( "A" , 0 , 1 ) ;
1581+
1582+ let netuid = Subtensor :: < T > :: get_next_netuid ( ) ;
1583+
1584+ let lock_cost = Subtensor :: < T > :: get_network_lock_cost ( ) ;
1585+ Subtensor :: < T > :: add_balance_to_coldkey_account ( & coldkey, lock_cost. into ( ) ) ;
1586+
1587+ assert_ok ! ( Subtensor :: <T >:: register_network(
1588+ RawOrigin :: Signed ( coldkey. clone( ) ) . into( ) ,
1589+ hotkey. clone( )
1590+ ) ) ;
1591+
1592+ SubtokenEnabled :: < T > :: insert ( netuid, true ) ;
1593+ Subtensor :: < T > :: set_network_pow_registration_allowed ( netuid, true ) ;
1594+ NetworkRegistrationAllowed :: < T > :: insert ( netuid, true ) ;
1595+ FirstEmissionBlockNumber :: < T > :: insert ( netuid, 0 ) ;
1596+
1597+ SubnetMechanism :: < T > :: insert ( netuid, 1 ) ;
1598+ SubnetworkN :: < T > :: insert ( netuid, 1 ) ;
1599+ Subtensor :: < T > :: set_tao_weight ( u64:: MAX ) ; // Set TAO weight to 1.0
1600+
1601+ let root_stake = 100_000_000u64 ;
1602+ Subtensor :: < T > :: increase_stake_for_hotkey_and_coldkey_on_subnet (
1603+ & hotkey,
1604+ & coldkey,
1605+ NetUid :: ROOT ,
1606+ root_stake. into ( ) ,
1607+ ) ;
1608+
1609+ let initial_total_hotkey_alpha = 100_000_000u64 ;
1610+ Subtensor :: < T > :: increase_stake_for_hotkey_and_coldkey_on_subnet (
1611+ & hotkey,
1612+ & coldkey,
1613+ netuid,
1614+ initial_total_hotkey_alpha. into ( ) ,
1615+ ) ;
1616+
1617+ let pending_root_alpha = 10_000_000u64 ;
1618+ Subtensor :: < T > :: drain_pending_emission (
1619+ netuid,
1620+ AlphaCurrency :: ZERO ,
1621+ pending_root_alpha. into ( ) ,
1622+ AlphaCurrency :: ZERO ,
1623+ ) ;
1624+
1625+ let initial_stake =
1626+ Subtensor :: < T > :: get_stake_for_hotkey_and_coldkey_on_subnet ( & hotkey, & coldkey, netuid) ;
1627+
1628+ assert_ok ! ( Subtensor :: <T >:: set_root_claim_type(
1629+ RawOrigin :: Signed ( coldkey. clone( ) ) . into( ) ,
1630+ RootClaimTypeEnum :: Keep
1631+ ) , ) ;
1632+
1633+ #[ extrinsic_call]
1634+ _ ( RawOrigin :: Signed ( coldkey. clone ( ) ) , BTreeSet :: from ( [ netuid] ) ) ;
1635+
1636+ // Verification
1637+ let new_stake =
1638+ Subtensor :: < T > :: get_stake_for_hotkey_and_coldkey_on_subnet ( & hotkey, & coldkey, netuid) ;
1639+
1640+ assert ! ( new_stake > initial_stake) ;
1641+ }
1642+
1643+ #[ benchmark]
1644+ fn sudo_set_num_root_claims ( ) {
1645+ #[ extrinsic_call]
1646+ _ ( RawOrigin :: Root , 40 ) ;
1647+ }
1648+
1649+ #[ benchmark]
1650+ fn sudo_set_root_claim_threshold ( ) {
1651+ let coldkey: T :: AccountId = whitelisted_caller ( ) ;
1652+ let hotkey: T :: AccountId = account ( "A" , 0 , 1 ) ;
1653+
1654+ let netuid = Subtensor :: < T > :: get_next_netuid ( ) ;
1655+
1656+ let lock_cost = Subtensor :: < T > :: get_network_lock_cost ( ) ;
1657+ Subtensor :: < T > :: add_balance_to_coldkey_account ( & coldkey, lock_cost. into ( ) ) ;
1658+
1659+ assert_ok ! ( Subtensor :: <T >:: register_network(
1660+ RawOrigin :: Signed ( coldkey. clone( ) ) . into( ) ,
1661+ hotkey. clone( )
1662+ ) ) ;
1663+
1664+ #[ extrinsic_call]
1665+ _ ( RawOrigin :: Root , netuid, 100 ) ;
1666+ }
15681667}
0 commit comments