Skip to content

Commit

Permalink
Merge pull request #614 from opentensor/feat/cold_key_swap
Browse files Browse the repository at this point in the history 
adds tests fixes remove hotkey
  • Loading branch information
@garrett-opentensor
garrett-opentensor committed Jul 5, 2024
2 parents a3ef6e0 + 5a26348 commit a62975c
Show file tree
Hide file tree
Showing 5 changed files with 70 additions and 80 deletions.
4 changes: 0 additions & 4 deletions pallets/subtensor/src/events.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -145,10 +145,6 @@ mod events {
current_coldkey: T::AccountId,
/// The account ID of the new coldkey
new_coldkey: T::AccountId,
/// The account ID of the hotkey
hotkey: T::AccountId,
/// The current stake of the hotkey
current_stake: u64,
/// The total balance of the hotkey
total_balance: <<T as Config>::Currency as fungible::Inspect<
<T as frame_system::Config>::AccountId,
Expand Down
3 changes: 1 addition & 2 deletions pallets/subtensor/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2030,11 +2030,10 @@ pub mod pallet {
.saturating_add(T::DbWeight::get().writes(527)), DispatchClass::Operational, Pays::No))]
pub fn unstake_all_and_transfer_to_new_coldkey(
origin: OriginFor<T>,
hotkey: T::AccountId,
new_coldkey: T::AccountId,
) -> DispatchResult {
let current_coldkey = ensure_signed(origin)?;
Self::do_unstake_all_and_transfer_to_new_coldkey(current_coldkey, hotkey, new_coldkey)
Self::do_unstake_all_and_transfer_to_new_coldkey(current_coldkey, new_coldkey)
}

// ---- SUDO ONLY FUNCTIONS ------------------------------------------------------------
Expand Down
51 changes: 25 additions & 26 deletions pallets/subtensor/src/staking.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -842,7 +842,6 @@ impl<T: Config> Pallet<T> {
/// # Arguments
///
/// * `current_coldkey` - The AccountId of the current coldkey.
/// * `hotkey` - The AccountId of the hotkey whose balance is being unstaked and transferred.
/// * `new_coldkey` - The AccountId of the new coldkey to receive the unstaked tokens.
///
/// # Returns
Expand All @@ -865,36 +864,38 @@ impl<T: Config> Pallet<T> {
///
pub fn do_unstake_all_and_transfer_to_new_coldkey(
current_coldkey: T::AccountId,
hotkey: T::AccountId,
new_coldkey: T::AccountId,
) -> DispatchResult {
// Ensure the hotkey exists and is owned by the current coldkey
ensure!(
Self::hotkey_account_exists(&hotkey),
Error::<T>::HotKeyAccountNotExists
);
ensure!(
Self::coldkey_owns_hotkey(&current_coldkey, &hotkey),
Error::<T>::NonAssociatedColdKey
);

// Ensure the new coldkey is different from the current one
ensure!(current_coldkey != new_coldkey, Error::<T>::SameColdkey);

// Get the current stake
let current_stake: u64 = Self::get_stake_for_coldkey_and_hotkey(&current_coldkey, &hotkey);
// Get all the hotkeys associated with this coldkey
let hotkeys: Vec<T::AccountId> = OwnedHotkeys::<T>::get(&current_coldkey);

// Unstake all balance if there's any stake
if current_stake > 0 {
Self::do_remove_stake(
RawOrigin::Signed(current_coldkey.clone()).into(),
hotkey.clone(),
current_stake,
)?;
}
// iterate over all hotkeys.
for next_hotkey in hotkeys {
ensure!(
Self::hotkey_account_exists(&next_hotkey),
Error::<T>::HotKeyAccountNotExists
);
ensure!(
Self::coldkey_owns_hotkey(&current_coldkey, &next_hotkey),
Error::<T>::NonAssociatedColdKey
);

// Get the total balance of the current coldkey account
// let total_balance: <<T as Config>::Currency as fungible::Inspect<<T as system::Config>::AccountId>>::Balance = T::Currency::total_balance(&current_coldkey);
// Get the current stake
let current_stake: u64 =
Self::get_stake_for_coldkey_and_hotkey(&current_coldkey, &next_hotkey);

// Unstake all balance if there's any stake
if current_stake > 0 {
Self::do_remove_stake(
RawOrigin::Signed(current_coldkey.clone()).into(),
next_hotkey.clone(),
current_stake,
)?;
}
}

let total_balance = Self::get_coldkey_balance(&current_coldkey);
log::info!("Total Bank Balance: {:?}", total_balance);
Expand All @@ -914,8 +915,6 @@ impl<T: Config> Pallet<T> {
Self::deposit_event(Event::AllBalanceUnstakedAndTransferredToNewColdkey {
current_coldkey: current_coldkey.clone(),
new_coldkey: new_coldkey.clone(),
hotkey: hotkey.clone(),
current_stake,
total_balance,
});

Expand Down
90 changes: 43 additions & 47 deletions pallets/subtensor/tests/staking.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3160,7 +3160,6 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_success() {

assert_ok!(SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey(
current_coldkey,
hotkey,
new_coldkey
));

Expand All @@ -3175,59 +3174,21 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_success() {
Event::AllBalanceUnstakedAndTransferredToNewColdkey {
current_coldkey,
new_coldkey,
hotkey,
current_stake: 500,
total_balance: 1000,
}
.into(),
);
});
}

#[test]
fn test_do_unstake_all_and_transfer_to_new_coldkey_hotkey_not_exists() {
new_test_ext(1).execute_with(|| {
let current_coldkey = U256::from(1);
let hotkey = U256::from(2);
let new_coldkey = U256::from(3);

assert_err!(
SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey(
current_coldkey,
hotkey,
new_coldkey
),
Error::<Test>::HotKeyAccountNotExists
);
});
}

#[test]
fn test_do_unstake_all_and_transfer_to_new_coldkey_non_associated_coldkey() {
new_test_ext(1).execute_with(|| {
let (_, hotkey, new_coldkey) = setup_test_environment();
let wrong_coldkey = U256::from(4);

assert_noop!(
SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey(
wrong_coldkey,
hotkey,
new_coldkey
),
Error::<Test>::NonAssociatedColdKey
);
});
}

#[test]
fn test_do_unstake_all_and_transfer_to_new_coldkey_same_coldkey() {
new_test_ext(1).execute_with(|| {
let (current_coldkey, hotkey, _) = setup_test_environment();
let (current_coldkey, _hotkey, _) = setup_test_environment();

assert_noop!(
SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey(
current_coldkey,
hotkey,
current_coldkey
),
Error::<Test>::SameColdkey
Expand Down Expand Up @@ -3270,7 +3231,6 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_no_balance() {
// Try to unstake and transfer
let result = SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey(
current_coldkey,
hotkey,
new_coldkey,
);

Expand Down Expand Up @@ -3343,7 +3303,6 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_with_no_stake() {
// Perform unstake and transfer
assert_ok!(SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey(
current_coldkey,
hotkey,
new_coldkey
));

Expand All @@ -3370,8 +3329,6 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_with_no_stake() {
Event::AllBalanceUnstakedAndTransferredToNewColdkey {
current_coldkey,
new_coldkey,
hotkey,
current_stake: 0,
total_balance: initial_balance,
}
.into(),
Expand All @@ -3394,7 +3351,6 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_with_multiple_stakes() {

assert_ok!(SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey(
current_coldkey,
hotkey,
new_coldkey
));

Expand All @@ -3409,8 +3365,48 @@ fn test_do_unstake_all_and_transfer_to_new_coldkey_with_multiple_stakes() {
Event::AllBalanceUnstakedAndTransferredToNewColdkey {
current_coldkey,
new_coldkey,
hotkey,
current_stake: 800,
total_balance: 1000,
}
.into(),
);
});
}

#[test]
fn test_do_unstake_all_and_transfer_to_new_coldkey_with_multiple_stakes_multiple() {
new_test_ext(1).execute_with(|| {
// Register the neuron to a new network
let netuid = 1;
let hotkey0 = U256::from(1);
let hotkey2 = U256::from(2);
let current_coldkey = U256::from(3);
let new_coldkey = U256::from(4);
add_network(netuid, 0, 0);
register_ok_neuron(1, hotkey0, current_coldkey, 0);
register_ok_neuron(1, hotkey2, current_coldkey, 0);
SubtensorModule::set_target_stakes_per_interval(10);
SubtensorModule::add_balance_to_coldkey_account(&current_coldkey, 1000);
assert_ok!(SubtensorModule::add_stake(
RuntimeOrigin::signed(current_coldkey),
hotkey0,
500
));
assert_ok!(SubtensorModule::add_stake(
RuntimeOrigin::signed(current_coldkey),
hotkey2,
300
));
assert_ok!(SubtensorModule::do_unstake_all_and_transfer_to_new_coldkey(
current_coldkey,
new_coldkey
));
assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey0), 0);
assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey2), 0);
assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), 1000);
System::assert_last_event(
Event::AllBalanceUnstakedAndTransferredToNewColdkey {
current_coldkey,
new_coldkey,
total_balance: 1000,
}
.into(),
Expand Down
2 changes: 1 addition & 1 deletion runtime/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
// `spec_version`, and `authoring_version` are the same between Wasm and native.
// This value is set to 100 to notify Polkadot-JS App (https://polkadot.js.org/apps) to use
// the compatible custom types.
spec_version: 156,
spec_version: 158,
impl_version: 1,
apis: RUNTIME_API_VERSIONS,
transaction_version: 1,
Expand Down

0 comments on commit a62975c

Please sign in to comment.