Skip to content

Commit af555d7

Browse files
author
Samuel Dare
committed
chore: introduce difficulty var to make tests shorter , fix test_comprehensive_coldkey_swap_scenarios
1 parent e4c9ef2 commit af555d7

File tree

6 files changed

+356
-180
lines changed

6 files changed

+356
-180
lines changed

‎pallets/subtensor/src/lib.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ pub mod pallet {
8787
/// Minimum balance required to perform a coldkey swap
8888
pub const MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP: u64 = 1_000_000_000; // 1 TAO in RAO
8989

90+
9091
#[pallet::pallet]
9192
#[pallet::without_storage_info]
9293
#[pallet::storage_version(STORAGE_VERSION)]
@@ -255,6 +256,9 @@ pub mod pallet {
255256
/// A flag to indicate if Liquid Alpha is enabled.
256257
#[pallet::constant]
257258
type LiquidAlphaOn: Get<bool>;
259+
/// The base difficulty for proof of work for coldkey swaps
260+
#[pallet::constant]
261+
type InitialBaseDifficulty: Get<u64>;
258262
}
259263

260264
/// Alias for the account ID.
@@ -331,6 +335,12 @@ pub mod pallet {
331335
360
332336
}
333337

338+
/// Default base difficulty for proof of work for coldkey swaps
339+
#[pallet::type_value]
340+
pub fn DefaultBaseDifficulty<T: Config>() -> u64 {
341+
T::InitialBaseDifficulty::get()
342+
}
343+
334344
#[pallet::storage] // --- ITEM ( total_stake )
335345
pub type TotalStake<T> = StorageValue<_, u64, ValueQuery>;
336346
#[pallet::storage] // --- ITEM ( default_take )
@@ -344,6 +354,9 @@ pub mod pallet {
344354
#[pallet::storage] // --- ITEM (target_stakes_per_interval)
345355
pub type TargetStakesPerInterval<T> =
346356
StorageValue<_, u64, ValueQuery, DefaultTargetStakesPerInterval<T>>;
357+
358+
#[pallet::storage] // --- ITEM ( base_difficulty )
359+
pub type BaseDifficulty<T> = StorageValue<_, u64, ValueQuery, DefaultBaseDifficulty<T>>;
347360
#[pallet::storage] // --- ITEM (default_stake_interval)
348361
pub type StakeInterval<T> = StorageValue<_, u64, ValueQuery, DefaultStakeInterval<T>>;
349362
#[pallet::storage] // --- MAP ( hot ) --> stake | Returns the total amount of stake under a hotkey.

‎pallets/subtensor/src/swap.rs‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,7 @@ impl<T: Config> Pallet<T> {
199199
}
200200
}
201201

202-
pub fn meets_min_allowed_coldkey_balance( coldkey: &T::AccountId ) -> bool {
202+
pub fn meets_min_allowed_coldkey_balance(coldkey: &T::AccountId) -> bool {
203203
let all_staked_keys: Vec<T::AccountId> = StakingHotkeys::<T>::get(coldkey);
204204
let mut total_staking_balance: u64 = 0;
205205
for hotkey in all_staked_keys {
@@ -225,12 +225,12 @@ impl<T: Config> Pallet<T> {
225225
///
226226
/// # Errors
227227
///
228-
/// This function will return an error if:
229-
/// - The old coldkey is the same as the new coldkey.
230-
/// - The new coldkey is already in the list of destination coldkeys.
231-
/// - There are already 2 destination coldkeys for the old coldkey.
232-
/// - The old coldkey doesn't have the minimum required TAO balance.
233-
/// - The proof of work is invalid or doesn't meet the required difficulty.
228+
229+
/// - `SameColdkey`: The old coldkey is the same as the new coldkey.
230+
/// - `DuplicateColdkey`: The new coldkey is already in the list of destination coldkeys.
231+
/// - `MaxColdkeyDestinationsReached`: There are already the maximum allowed destination coldkeys for the old coldkey.
232+
/// - `InsufficientBalanceToPerformColdkeySwap`: The old coldkey doesn't have the minimum required TAO balance.
233+
/// - `InvalidDifficulty`: The proof of work is invalid or doesn't meet the required difficulty.
234234
///
235235
/// # Notes
236236
///
@@ -319,7 +319,7 @@ impl<T: Config> Pallet<T> {
319319
/// Calculate the proof of work difficulty based on the number of swap attempts
320320
#[allow(clippy::arithmetic_side_effects)]
321321
pub fn calculate_pow_difficulty(swap_attempts: u32) -> U256 {
322-
let base_difficulty: U256 = U256::from(10_000_000); // Base difficulty
322+
let base_difficulty: U256 = U256::from(BaseDifficulty::<T>::get()); // Base difficulty
323323
base_difficulty.saturating_mul(U256::from(2).pow(U256::from(swap_attempts)))
324324
}
325325

‎pallets/subtensor/tests/mock.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,7 @@ parameter_types! {
168168
pub const InitialAlphaHigh: u16 = 58982; // Represents 0.9 as per the production default
169169
pub const InitialAlphaLow: u16 = 45875; // Represents 0.7 as per the production default
170170
pub const InitialLiquidAlphaOn: bool = false; // Default value for LiquidAlphaOn
171+
pub const SubtensorInitialBaseDifficulty: u64 = 10_000; // Base difficulty
171172
}
172173

173174
// Configure collective pallet for council
@@ -372,6 +373,7 @@ impl pallet_subtensor::Config for Test {
372373
type AlphaHigh = InitialAlphaHigh;
373374
type AlphaLow = InitialAlphaLow;
374375
type LiquidAlphaOn = InitialLiquidAlphaOn;
376+
type InitialBaseDifficulty = SubtensorInitialBaseDifficulty;
375377
}
376378

377379
impl pallet_utility::Config for Test {

‎pallets/subtensor/tests/staking.rs‎

Lines changed: 81 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -3167,8 +3167,11 @@ fn test_arbitrated_coldkey_swap_success() {
31673167
let (current_coldkey, hotkey, new_coldkey) = setup_test_environment();
31683168

31693169
let current_block = SubtensorModule::get_current_block_as_u64();
3170-
let (work, nonce) =
3171-
generate_valid_pow(&current_coldkey, current_block, U256::from(10_000_000u64));
3170+
let (work, nonce) = generate_valid_pow(
3171+
&current_coldkey,
3172+
current_block,
3173+
U256::from(BaseDifficulty::<Test>::get()),
3174+
);
31723175
SubtensorModule::add_balance_to_coldkey_account(
31733176
&current_coldkey,
31743177
MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP,
@@ -3239,8 +3242,11 @@ fn test_arbitrated_coldkey_swap_same_coldkey() {
32393242
let (current_coldkey, _hotkey, _) = setup_test_environment();
32403243

32413244
let current_block = SubtensorModule::get_current_block_as_u64();
3242-
let (work, nonce) =
3243-
generate_valid_pow(&current_coldkey, current_block, U256::from(10_000_000u64));
3245+
let (work, nonce) = generate_valid_pow(
3246+
&current_coldkey,
3247+
current_block,
3248+
U256::from(BaseDifficulty::<Test>::get()),
3249+
);
32443250

32453251
assert_noop!(
32463252
SubtensorModule::do_schedule_coldkey_swap(
@@ -3290,8 +3296,11 @@ fn test_arbitrated_coldkey_swap_no_balance() {
32903296

32913297
// Generate valid PoW
32923298
let current_block = SubtensorModule::get_current_block_as_u64();
3293-
let (work, nonce) =
3294-
generate_valid_pow(&current_coldkey, current_block, U256::from(10_000_000u64));
3299+
let (work, nonce) = generate_valid_pow(
3300+
&current_coldkey,
3301+
current_block,
3302+
U256::from(BaseDifficulty::<Test>::get()),
3303+
);
32953304

32963305
// Try to schedule coldkey swap
32973306
let result = SubtensorModule::do_schedule_coldkey_swap(
@@ -3373,8 +3382,11 @@ fn test_arbitrated_coldkey_swap_with_no_stake() {
33733382
assert_eq!(Balances::total_balance(&new_coldkey), 0);
33743383

33753384
let current_block = SubtensorModule::get_current_block_as_u64();
3376-
let (work, nonce) =
3377-
generate_valid_pow(&current_coldkey, current_block, U256::from(10_000_000u64));
3385+
let (work, nonce) = generate_valid_pow(
3386+
&current_coldkey,
3387+
current_block,
3388+
U256::from(BaseDifficulty::<Test>::get()),
3389+
);
33783390

33793391
// Schedule coldkey swap
33803392
assert_ok!(SubtensorModule::do_schedule_coldkey_swap(
@@ -3435,8 +3447,11 @@ fn test_arbitrated_coldkey_swap_with_multiple_stakes() {
34353447
));
34363448

34373449
let current_block = SubtensorModule::get_current_block_as_u64();
3438-
let (work, nonce) =
3439-
generate_valid_pow(&current_coldkey, current_block, U256::from(10_000_000u64));
3450+
let (work, nonce) = generate_valid_pow(
3451+
&current_coldkey,
3452+
current_block,
3453+
U256::from(BaseDifficulty::<Test>::get()),
3454+
);
34403455

34413456
assert_ok!(SubtensorModule::do_schedule_coldkey_swap(
34423457
&current_coldkey.clone(),
@@ -3476,6 +3491,9 @@ fn test_arbitrated_coldkey_swap_with_multiple_stakes() {
34763491
#[test]
34773492
fn test_arbitrated_coldkey_swap_multiple_arbitrations() {
34783493
new_test_ext(1).execute_with(|| {
3494+
// Set a very low base difficulty for testing
3495+
BaseDifficulty::<Test>::put(1);
3496+
34793497
// Create coldkey with three choices.
34803498
let coldkey: AccountId = U256::from(1);
34813499
let new_coldkey1: AccountId = U256::from(2);
@@ -3493,12 +3511,11 @@ fn test_arbitrated_coldkey_swap_multiple_arbitrations() {
34933511
register_ok_neuron(1, hotkey, coldkey, 0);
34943512

34953513
let current_block = SubtensorModule::get_current_block_as_u64();
3496-
let (work1, nonce1) =
3497-
generate_valid_pow(&coldkey, current_block, U256::from(10_000_000u64));
3498-
let (work2, nonce2) =
3499-
generate_valid_pow(&coldkey, current_block, U256::from(20_000_000u64));
3500-
let (work3, nonce3) =
3501-
generate_valid_pow(&coldkey, current_block, U256::from(30_000_000u64));
3514+
3515+
// Generate valid PoW for each swap attempt
3516+
let (work1, nonce1) = generate_valid_pow(&coldkey, current_block, U256::from(1));
3517+
let (work2, nonce2) = generate_valid_pow(&coldkey, current_block, U256::from(2));
3518+
let (work3, nonce3) = generate_valid_pow(&coldkey, current_block, U256::from(4));
35023519

35033520
// Schedule three swaps
35043521
assert_ok!(SubtensorModule::do_schedule_coldkey_swap(
@@ -3531,7 +3548,6 @@ fn test_arbitrated_coldkey_swap_multiple_arbitrations() {
35313548

35323549
// Simulate the passage of blocks and on_idle calls
35333550
for i in 0..(5400 + 1) {
3534-
// Simulate 10 blocks
35353551
next_block();
35363552
SubtensorModule::on_idle(System::block_number(), Weight::MAX);
35373553

@@ -3655,8 +3671,11 @@ fn test_arbitration_period_extension() {
36553671
let another_coldkey = U256::from(4);
36563672

36573673
let current_block = SubtensorModule::get_current_block_as_u64();
3658-
let (work1, nonce1) =
3659-
generate_valid_pow(&current_coldkey, current_block, U256::from(10_000_000u64));
3674+
let (work1, nonce1) = generate_valid_pow(
3675+
&current_coldkey,
3676+
current_block,
3677+
U256::from(BaseDifficulty::<Test>::get()),
3678+
);
36603679
let (work2, nonce2) =
36613680
generate_valid_pow(&current_coldkey, current_block, U256::from(20_000_000u64));
36623681
SubtensorModule::add_balance_to_coldkey_account(
@@ -3725,10 +3744,16 @@ fn test_concurrent_arbitrated_coldkey_swaps() {
37253744
);
37263745

37273746
let current_block = SubtensorModule::get_current_block_as_u64();
3728-
let (work1, nonce1) =
3729-
generate_valid_pow(&coldkey1, current_block, U256::from(10_000_000u64));
3730-
let (work2, nonce2) =
3731-
generate_valid_pow(&coldkey2, current_block, U256::from(10_000_000u64));
3747+
let (work1, nonce1) = generate_valid_pow(
3748+
&coldkey1,
3749+
current_block,
3750+
U256::from(BaseDifficulty::<Test>::get()),
3751+
);
3752+
let (work2, nonce2) = generate_valid_pow(
3753+
&coldkey2,
3754+
current_block,
3755+
U256::from(BaseDifficulty::<Test>::get()),
3756+
);
37323757
// Schedule swaps for both coldkeys
37333758
assert_ok!(SubtensorModule::do_schedule_coldkey_swap(
37343759
&coldkey1.clone(),
@@ -3773,7 +3798,7 @@ fn test_get_remaining_arbitration_period() {
37733798
let (work, nonce) = generate_valid_pow(
37743799
&coldkey_account_id,
37753800
current_block,
3776-
U256::from(10_000_000u64),
3801+
U256::from(BaseDifficulty::<Test>::get()),
37773802
);
37783803

37793804
SubtensorModule::add_balance_to_coldkey_account(
@@ -3833,7 +3858,7 @@ fn test_transfer_coldkey_in_arbitration() {
38333858
let (work, nonce) = generate_valid_pow(
38343859
&coldkey_account_id,
38353860
current_block,
3836-
U256::from(10_000_000u64),
3861+
U256::from(BaseDifficulty::<Test>::get()),
38373862
);
38383863

38393864
// Schedule a coldkey swap to put the coldkey in arbitration
@@ -3876,8 +3901,11 @@ fn test_add_stake_coldkey_in_arbitration() {
38763901
);
38773902

38783903
let current_block = SubtensorModule::get_current_block_as_u64();
3879-
let (work, nonce) =
3880-
generate_valid_pow(&coldkey_account_id, current_block, U256::from(1_000_000u64));
3904+
let (work, nonce) = generate_valid_pow(
3905+
&coldkey_account_id,
3906+
current_block,
3907+
U256::from(BaseDifficulty::<Test>::get()),
3908+
);
38813909

38823910
// Schedule a coldkey swap to put the coldkey in arbitration
38833911
assert_ok!(SubtensorModule::do_schedule_coldkey_swap(
@@ -3916,8 +3944,11 @@ fn test_remove_stake_coldkey_in_arbitration() {
39163944
SubtensorModule::increase_stake_on_hotkey_account(&hotkey_account_id, 1000);
39173945

39183946
let current_block = SubtensorModule::get_current_block_as_u64();
3919-
let (work, nonce) =
3920-
generate_valid_pow(&coldkey_account_id, current_block, U256::from(1_000_000u64));
3947+
let (work, nonce) = generate_valid_pow(
3948+
&coldkey_account_id,
3949+
current_block,
3950+
U256::from(BaseDifficulty::<Test>::get()),
3951+
);
39213952

39223953
// Schedule a coldkey swap to put the coldkey in arbitration
39233954
assert_ok!(SubtensorModule::do_schedule_coldkey_swap(
@@ -3992,8 +4023,11 @@ fn test_coldkey_meets_enough() {
39924023
add_network(netuid, 13, 0);
39934024
register_ok_neuron(netuid, hotkey, coldkey, 0);
39944025
let current_block = SubtensorModule::get_current_block_as_u64();
3995-
let (work1, nonce1) =
3996-
generate_valid_pow(&coldkey, current_block, U256::from(10_000_000u64));
4026+
let (work1, nonce1) = generate_valid_pow(
4027+
&coldkey,
4028+
current_block,
4029+
U256::from(BaseDifficulty::<Test>::get()),
4030+
);
39974031
assert_err!(
39984032
SubtensorModule::do_schedule_coldkey_swap(
39994033
&coldkey.clone(),
@@ -4039,21 +4073,29 @@ fn test_comprehensive_coldkey_swap_scenarios() {
40394073
SubnetOwner::<Test>::insert(netuid1, subnet_owner1);
40404074
SubnetOwner::<Test>::insert(netuid2, subnet_owner2);
40414075

4042-
// Add balance to regular user
4076+
// Add balance to subnet owners and regular user
4077+
SubtensorModule::add_balance_to_coldkey_account(
4078+
&subnet_owner1,
4079+
MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP,
4080+
);
4081+
SubtensorModule::add_balance_to_coldkey_account(
4082+
&subnet_owner2,
4083+
MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP,
4084+
);
40434085
SubtensorModule::add_balance_to_coldkey_account(
40444086
&regular_user,
40454087
MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP * 2,
40464088
);
40474089

4090+
// Set a very low base difficulty for testing
4091+
BaseDifficulty::<Test>::put(1);
4092+
40484093
let current_block = SubtensorModule::get_current_block_as_u64();
40494094

40504095
// Schedule swaps for subnet owners and regular user
4051-
let (work1, nonce1) =
4052-
generate_valid_pow(&subnet_owner1, current_block, U256::from(10_000_000u64));
4053-
let (work2, nonce2) =
4054-
generate_valid_pow(&subnet_owner2, current_block, U256::from(20_000_000u64));
4055-
let (work3, nonce3) =
4056-
generate_valid_pow(&regular_user, current_block, U256::from(30_000_000u64));
4096+
let (work1, nonce1) = generate_valid_pow(&subnet_owner1, current_block, U256::from(BaseDifficulty::<Test>::get()));
4097+
let (work2, nonce2) = generate_valid_pow(&subnet_owner2, current_block, U256::from(2) * U256::from(BaseDifficulty::<Test>::get()));
4098+
let (work3, nonce3) = generate_valid_pow(&regular_user, current_block, U256::from(3)* U256::from(BaseDifficulty::<Test>::get()),);
40574099

40584100
assert_ok!(SubtensorModule::do_schedule_coldkey_swap(
40594101
&subnet_owner1,
@@ -4110,7 +4152,7 @@ fn test_comprehensive_coldkey_swap_scenarios() {
41104152
let (work4, nonce4) = generate_valid_pow(
41114153
&subnet_owner1,
41124154
current_block + i as u64,
4113-
U256::from(40_000_000u64),
4155+
U256::from(4) * U256::from(BaseDifficulty::<Test>::get()),
41144156
);
41154157
assert_ok!(SubtensorModule::do_schedule_coldkey_swap(
41164158
&subnet_owner1,
@@ -4151,54 +4193,5 @@ fn test_comprehensive_coldkey_swap_scenarios() {
41514193
MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP * 2
41524194
);
41534195
assert_eq!(SubtensorModule::get_coldkey_balance(&regular_user), 0);
4154-
4155-
// Test multiple calls for the same subnet owner
4156-
let (work5, nonce5) =
4157-
generate_valid_pow(&new_coldkey1, current_block + 7, U256::from(50_000_000u64));
4158-
let (work6, nonce6) =
4159-
generate_valid_pow(&new_coldkey1, current_block + 7, U256::from(60_000_000u64));
4160-
4161-
assert_ok!(SubtensorModule::do_schedule_coldkey_swap(
4162-
&new_coldkey1,
4163-
&subnet_owner1,
4164-
work5.to_fixed_bytes().to_vec(),
4165-
current_block + 7,
4166-
nonce5
4167-
));
4168-
4169-
assert_ok!(SubtensorModule::do_schedule_coldkey_swap(
4170-
&new_coldkey1,
4171-
&subnet_owner2,
4172-
work6.to_fixed_bytes().to_vec(),
4173-
current_block + 7,
4174-
nonce6
4175-
));
4176-
4177-
assert_eq!(
4178-
ColdkeySwapDestinations::<Test>::get(new_coldkey1),
4179-
vec![subnet_owner1, subnet_owner2]
4180-
);
4181-
4182-
// Run through another arbitration period plus one block
4183-
for i in 0..6 {
4184-
next_block();
4185-
SubtensorModule::on_idle(System::block_number(), Weight::MAX);
4186-
4187-
log::info!(
4188-
"Block {}: Coldkey in arbitration: {}, Swap destinations: {:?}",
4189-
i + 7,
4190-
SubtensorModule::coldkey_in_arbitration(&new_coldkey1),
4191-
ColdkeySwapDestinations::<Test>::get(new_coldkey1)
4192-
);
4193-
}
4194-
4195-
// Check final state
4196-
log::info!(
4197-
"Final state - Swap destinations for new_coldkey1: {:?}",
4198-
ColdkeySwapDestinations::<Test>::get(new_coldkey1)
4199-
);
4200-
assert!(ColdkeySwapDestinations::<Test>::get(new_coldkey1).is_empty());
4201-
assert_eq!(SubnetOwner::<Test>::get(netuid1), subnet_owner1);
4202-
assert_eq!(SubnetOwner::<Test>::get(netuid2), subnet_owner2);
42034196
});
42044197
}

0 commit comments

Comments
 (0)