Skip to content

Commit

Permalink
test: different staking amounts
Browse files Browse the repository at this point in the history
  • Loading branch information
Samuel Dare committed Jul 15, 2024
1 parent cd47404 commit 7f9afbe
Showing 1 changed file with 139 additions and 0 deletions.
139 changes: 139 additions & 0 deletions pallets/subtensor/tests/swap.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2114,3 +2114,142 @@ fn test_coldkey_swap_with_hotfix_and_total_stake_fix() {
assert_eq!(Stake::<Test>::get(h3, old_coldkey), 0);
});
}

// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test swap -- test_coldkey_swap_with_hotfix_and_total_stake_fix_different_amounts --exact --nocapture
#[test]
fn test_coldkey_swap_with_hotfix_and_total_stake_fix_different_amounts() {
new_test_ext(1).execute_with(|| {
let new_coldkey = U256::from(0);
let old_coldkey = U256::from(4);
let h1 = U256::from(5);
let h2 = U256::from(6);
let h3 = U256::from(7);
let stake_amount_1 = 100u64;
let stake_amount_2 = 150u64;
let stake_amount_3 = 200u64;

// Setup initial state to simulate after a broken coldkey swap
Stake::<Test>::insert(h3, old_coldkey, stake_amount_3);
Stake::<Test>::insert(h2, old_coldkey, stake_amount_2);
Stake::<Test>::insert(h1, new_coldkey, stake_amount_1);
assert_eq!(Stake::<Test>::get(h3, old_coldkey), stake_amount_3);
assert_eq!(Stake::<Test>::get(h2, old_coldkey), stake_amount_2);
assert_eq!(Stake::<Test>::get(h1, new_coldkey), stake_amount_1);

StakingHotkeys::<Test>::insert(new_coldkey, vec![h1, h2]);
StakingHotkeys::<Test>::insert(old_coldkey, vec![h3, h2]);

// Set up total coldkey stake (this would be incorrect after a broken swap)
TotalColdkeyStake::<Test>::insert(old_coldkey, stake_amount_2 + stake_amount_3);
TotalColdkeyStake::<Test>::insert(new_coldkey, stake_amount_1);

// Log initial state
log::info!("Initial state (simulating after broken swap):");
log::info!(
"Old coldkey stake: {}",
SubtensorModule::get_total_stake_for_coldkey(&old_coldkey)
);
log::info!(
"New coldkey stake: {}",
SubtensorModule::get_total_stake_for_coldkey(&new_coldkey)
);
log::info!(
"Old coldkey staking hotkeys: {:?}",
StakingHotkeys::<Test>::get(old_coldkey)
);
log::info!(
"New coldkey staking hotkeys: {:?}",
StakingHotkeys::<Test>::get(new_coldkey)
);
log::info!("h1 stake (new): {}", Stake::<Test>::get(h1, new_coldkey));
log::info!("h2 stake (old): {}", Stake::<Test>::get(h2, old_coldkey));
log::info!("h2 stake (new): {}", Stake::<Test>::get(h2, new_coldkey));
log::info!("h3 stake (old): {}", Stake::<Test>::get(h3, old_coldkey));

// Apply the hotfix for StakingHotkeys using the sudo extrinsic
assert_ok!(SubtensorModule::sudo_hotfix_swap_coldkey_delegates(
RuntimeOrigin::root(),
old_coldkey,
new_coldkey
));

// Log state after hotfix
log::info!("After hotfix:");
log::info!(
"Old coldkey stake: {}",
SubtensorModule::get_total_stake_for_coldkey(&old_coldkey)
);
log::info!(
"New coldkey stake: {}",
SubtensorModule::get_total_stake_for_coldkey(&new_coldkey)
);
log::info!(
"Old coldkey staking hotkeys: {:?}",
StakingHotkeys::<Test>::get(old_coldkey)
);
log::info!(
"New coldkey staking hotkeys: {:?}",
StakingHotkeys::<Test>::get(new_coldkey)
);
log::info!("h1 stake (new): {}", Stake::<Test>::get(h1, new_coldkey));
log::info!("h2 stake (old): {}", Stake::<Test>::get(h2, old_coldkey));
log::info!("h2 stake (new): {}", Stake::<Test>::get(h2, new_coldkey));
log::info!("h3 stake (new): {}", Stake::<Test>::get(h3, new_coldkey));

// Fix the TotalColdkeyStake using the sudo extrinsic
assert_ok!(SubtensorModule::sudo_hotfix_swap_coldkey_total_stake(
RuntimeOrigin::root(),
old_coldkey,
new_coldkey
));

// Log final state
log::info!("Final state:");
let final_old_coldkey_stake = SubtensorModule::get_total_stake_for_coldkey(&old_coldkey);
let final_new_coldkey_stake = SubtensorModule::get_total_stake_for_coldkey(&new_coldkey);
log::info!("Old coldkey stake: {}", final_old_coldkey_stake);
log::info!("New coldkey stake: {}", final_new_coldkey_stake);
log::info!(
"Old coldkey staking hotkeys: {:?}",
StakingHotkeys::<Test>::get(old_coldkey)
);
log::info!(
"New coldkey staking hotkeys: {:?}",
StakingHotkeys::<Test>::get(new_coldkey)
);
log::info!("h1 stake (new): {}", Stake::<Test>::get(h1, new_coldkey));
log::info!("h2 stake (old): {}", Stake::<Test>::get(h2, old_coldkey));
log::info!("h2 stake (new): {}", Stake::<Test>::get(h2, new_coldkey));
log::info!("h3 stake (new): {}", Stake::<Test>::get(h3, new_coldkey));

// Verify final state
assert_eq!(
final_old_coldkey_stake, 0,
"Expected old coldkey stake to be 0, but got {}",
final_old_coldkey_stake
);
assert_eq!(
final_new_coldkey_stake,
stake_amount_1 + stake_amount_2 + stake_amount_3,
"Expected total stake for new coldkey to be {}, but got {}",
stake_amount_1 + stake_amount_2 + stake_amount_3,
final_new_coldkey_stake
);
assert_eq!(StakingHotkeys::<Test>::get(old_coldkey), Vec::<U256>::new());
assert_eq!(StakingHotkeys::<Test>::get(new_coldkey), vec![h1, h2, h3]);

// Verify individual stakes
assert_eq!(Stake::<Test>::get(h1, new_coldkey), stake_amount_1);
assert_eq!(
Stake::<Test>::get(h2, new_coldkey),
stake_amount_2,
"Expected h2 stake to be {}, but got {}",
stake_amount_2,
Stake::<Test>::get(h2, new_coldkey)
);
assert_eq!(Stake::<Test>::get(h3, new_coldkey), stake_amount_3);
assert_eq!(Stake::<Test>::get(h1, old_coldkey), 0);
assert_eq!(Stake::<Test>::get(h2, old_coldkey), 0);
assert_eq!(Stake::<Test>::get(h3, old_coldkey), 0);
});
}

0 comments on commit 7f9afbe

Please sign in to comment.