diff --git a/pallets/subtensor/src/migrations/migrate_fix_pending_emission.rs b/pallets/subtensor/src/migrations/migrate_fix_pending_emission.rs new file mode 100644 index 0000000000..b5e833aeb9 --- /dev/null +++ b/pallets/subtensor/src/migrations/migrate_fix_pending_emission.rs @@ -0,0 +1,501 @@ +use super::*; +use alloc::string::String; +use frame_support::{traits::Get, weights::Weight}; +use sp_core::crypto::Ss58Codec; +use sp_runtime::AccountId32; + +fn get_account_id_from_ss58(ss58_str: &str) -> Result { + let account = + AccountId32::from_ss58check(ss58_str).map_err(|_| codec::Error::from("Invalid SS58"))?; + let onchain_account = T::AccountId::decode(&mut account.as_ref())?; + + Ok(onchain_account) +} + +/** + * Migrates the pending emissions from the old hotkey to the new hotkey. + * Also migrates the stake entry of (old_hotkey, 0x000) to the pending emissions of the new hotkey. + */ +fn migrate_pending_emissions_including_null_stake( + old_hotkey: &T::AccountId, + new_hotkey: &T::AccountId, + migration_account: &T::AccountId, +) -> Weight { + let mut weight = T::DbWeight::get().reads(0); + let null_account = &DefaultAccount::::get(); + weight.saturating_accrue(T::DbWeight::get().reads(1)); + + // Get the pending emissions for the OLD hotkey + let pending_emissions_old: u64 = PendingdHotkeyEmission::::get(old_hotkey); + PendingdHotkeyEmission::::remove(old_hotkey); + weight.saturating_accrue(T::DbWeight::get().reads(1)); + + // Get the stake for the 0x000 key + let null_stake = Stake::::get(old_hotkey, null_account); + weight.saturating_accrue(T::DbWeight::get().reads(1)); + // Remove + Stake::::remove(old_hotkey, null_account); + weight.saturating_accrue(T::DbWeight::get().writes(1)); + + let new_total_coldkey_stake = + TotalColdkeyStake::::get(null_account).saturating_sub(null_stake); + if new_total_coldkey_stake == 0 { + TotalColdkeyStake::::remove(null_account); + } else { + TotalColdkeyStake::::insert(null_account, new_total_coldkey_stake); + } + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1)); + + let new_staking_hotkeys = StakingHotkeys::::get(null_account); + let new_staking_hotkeys = new_staking_hotkeys + .into_iter() + .filter(|hk| hk != old_hotkey) + .collect::>(); + StakingHotkeys::::insert(null_account, new_staking_hotkeys); + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1)); + + // Insert the stake from the null account to the MIGRATION account under the OLD hotkey + Stake::::insert(old_hotkey, migration_account, null_stake); + TotalColdkeyStake::::insert( + migration_account, + TotalColdkeyStake::::get(migration_account).saturating_add(null_stake), + ); + let mut new_staking_hotkeys = StakingHotkeys::::get(migration_account); + if !new_staking_hotkeys.contains(old_hotkey) { + new_staking_hotkeys.push(old_hotkey.clone()); + } + StakingHotkeys::::insert(migration_account, new_staking_hotkeys); + weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 3)); + + // Get the pending emissions for the NEW hotkey + let pending_emissions_new: u64 = PendingdHotkeyEmission::::get(new_hotkey); + weight.saturating_accrue(T::DbWeight::get().reads(1)); + + // Add the pending emissions for the new hotkey and the old hotkey + PendingdHotkeyEmission::::insert( + new_hotkey, + pending_emissions_new.saturating_add(pending_emissions_old), + ); + weight.saturating_accrue(T::DbWeight::get().writes(1)); + + weight +} + +// This executes the migration to fix the pending emissions +// This also migrates the stake entry of (old_hotkey, 0x000) to the Migration Account for +// both the old hotkeys. +pub fn do_migrate_fix_pending_emission() -> Weight { + // Initialize the weight with one read operation. + let mut weight = T::DbWeight::get().reads(1); + + let taostats_old_hotkey = "5Hddm3iBFD2GLT5ik7LZnT3XJUnRnN8PoeCFgGQgawUVKNm8"; + let taostats_new_hotkey = "5GKH9FPPnWSUoeeTJp19wVtd84XqFW4pyK2ijV2GsFbhTrP1"; + let migration_coldkey = "5GeRjQYsobRWFnrbBmGe5ugme3rfnDVF69N45YtdBpUFsJG8"; + + let taostats_old_hk_account = get_account_id_from_ss58::(taostats_old_hotkey); + let taostats_new_hk_account = get_account_id_from_ss58::(taostats_new_hotkey); + let migration_ck_account = get_account_id_from_ss58::(migration_coldkey); + + match ( + taostats_old_hk_account, + taostats_new_hk_account, + migration_ck_account.clone(), + ) { + (Ok(taostats_old_hk_acct), Ok(taostats_new_hk_acct), Ok(migration_ck_account)) => { + weight.saturating_accrue(migrate_pending_emissions_including_null_stake::( + &taostats_old_hk_acct, + &taostats_new_hk_acct, + &migration_ck_account, + )); + log::info!("Migrated pending emissions from taostats old hotkey to new hotkey"); + } + _ => { + log::warn!("Failed to get account id from ss58 for taostats hotkeys"); + return weight; + } + } + + let datura_old_hotkey = "5FKstHjZkh4v3qAMSBa1oJcHCLjxYZ8SNTSz1opTv4hR7gVB"; + let datura_new_hotkey = "5GP7c3fFazW9GXK8Up3qgu2DJBk8inu4aK9TZy3RuoSWVCMi"; + + let datura_old_hk_account = get_account_id_from_ss58::(datura_old_hotkey); + let datura_new_hk_account = get_account_id_from_ss58::(datura_new_hotkey); + + match ( + datura_old_hk_account, + datura_new_hk_account, + migration_ck_account, + ) { + (Ok(datura_old_hk_acct), Ok(datura_new_hk_acct), Ok(migration_ck_account)) => { + weight.saturating_accrue(migrate_pending_emissions_including_null_stake::( + &datura_old_hk_acct, + &datura_new_hk_acct, + &migration_ck_account, + )); + log::info!("Migrated pending emissions from datura old hotkey to new hotkey"); + } + _ => { + log::warn!("Failed to get account id from ss58 for datura hotkeys"); + return weight; + } + } + + weight +} + +/// Collection of storage item formats from the previous storage version. +/// +/// Required so we can read values in the v0 storage format during the migration. +#[cfg(feature = "try-runtime")] +mod v0 { + use subtensor_macros::freeze_struct; + + #[freeze_struct("2228babfc0580c62")] + #[derive(codec::Encode, codec::Decode, Clone, PartialEq, Debug)] + pub struct OldStorage { + pub total_issuance_before: u64, + pub total_stake_before: u64, + pub expected_taostats_new_hk_pending_emission: u64, + pub expected_datura_new_hk_pending_emission: u64, + pub old_migration_stake_taostats: u64, + pub old_null_stake_taostats: u64, + pub old_migration_stake_datura: u64, + pub old_null_stake_datura: u64, + } +} + +impl Pallet { + #[cfg(feature = "try-runtime")] + fn check_null_stake_invariants( + old_storage: v0::OldStorage, + ) -> Result<(), sp_runtime::TryRuntimeError> { + let null_account = &DefaultAccount::::get(); + + let taostats_old_hotkey = "5Hddm3iBFD2GLT5ik7LZnT3XJUnRnN8PoeCFgGQgawUVKNm8"; + let taostats_new_hotkey = "5GKH9FPPnWSUoeeTJp19wVtd84XqFW4pyK2ijV2GsFbhTrP1"; + let migration_coldkey = "5GeRjQYsobRWFnrbBmGe5ugme3rfnDVF69N45YtdBpUFsJG8"; + + let taostats_old_hk_account = &get_account_id_from_ss58::(taostats_old_hotkey); + let taostats_new_hk_account = &get_account_id_from_ss58::(taostats_new_hotkey); + let migration_ck_account = &get_account_id_from_ss58::(migration_coldkey); + + let old = old_storage; + let null_stake_total = old + .old_null_stake_taostats + .saturating_add(old.old_null_stake_datura) + .saturating_add(old.old_migration_stake_taostats) + .saturating_add(old.old_migration_stake_datura); + + match ( + taostats_old_hk_account, + taostats_new_hk_account, + migration_ck_account, + ) { + (Ok(taostats_old_hk_acct), Ok(taostats_new_hk_acct), Ok(migration_ck_acct)) => { + // Check the pending emission is added to new the TaoStats hotkey + assert_eq!( + PendingdHotkeyEmission::::get(taostats_new_hk_acct), + old.expected_taostats_new_hk_pending_emission + ); + + assert_eq!(PendingdHotkeyEmission::::get(taostats_old_hk_acct), 0); + + assert_eq!(Stake::::get(taostats_old_hk_acct, null_account), 0); + + assert!(StakingHotkeys::::get(migration_ck_acct).contains(taostats_old_hk_acct)); + + assert_eq!( + Self::get_stake_for_coldkey_and_hotkey(null_account, taostats_old_hk_acct), + 0 + ); + + // Check the total hotkey stake is the same + assert_eq!( + TotalHotkeyStake::::get(taostats_old_hk_acct), + old.old_null_stake_taostats + .saturating_add(old.old_migration_stake_taostats) + ); + + let new_null_stake_taostats = + Self::get_stake_for_coldkey_and_hotkey(migration_ck_acct, taostats_old_hk_acct); + + assert_eq!( + new_null_stake_taostats, + old.old_null_stake_taostats + .saturating_add(old.old_migration_stake_taostats) + ); + } + _ => { + log::warn!("Failed to get account id from ss58 for taostats hotkeys"); + return Err("Failed to get account id from ss58 for taostats hotkeys".into()); + } + } + + let datura_old_hotkey = "5FKstHjZkh4v3qAMSBa1oJcHCLjxYZ8SNTSz1opTv4hR7gVB"; + let datura_new_hotkey = "5GP7c3fFazW9GXK8Up3qgu2DJBk8inu4aK9TZy3RuoSWVCMi"; + + let datura_old_hk_account = &get_account_id_from_ss58::(datura_old_hotkey); + let datura_new_hk_account = &get_account_id_from_ss58::(datura_new_hotkey); + + match ( + datura_old_hk_account, + datura_new_hk_account, + migration_ck_account, + ) { + (Ok(datura_old_hk_acct), Ok(datura_new_hk_acct), Ok(migration_ck_acct)) => { + // Check the pending emission is added to new Datura hotkey + assert_eq!( + crate::PendingdHotkeyEmission::::get(datura_new_hk_acct), + old.expected_datura_new_hk_pending_emission + ); + + // Check the pending emission is removed from old ones + assert_eq!(PendingdHotkeyEmission::::get(datura_old_hk_acct), 0); + + // Check the stake entry is removed + assert_eq!(Stake::::get(datura_old_hk_acct, null_account), 0); + + assert!(StakingHotkeys::::get(migration_ck_acct).contains(datura_old_hk_acct)); + + assert_eq!( + Self::get_stake_for_coldkey_and_hotkey(null_account, datura_old_hk_acct), + 0 + ); + + // Check the total hotkey stake is the same + assert_eq!( + TotalHotkeyStake::::get(datura_old_hk_acct), + old.old_null_stake_datura + .saturating_add(old.old_migration_stake_datura) + ); + + let new_null_stake_datura = + Self::get_stake_for_coldkey_and_hotkey(migration_ck_acct, datura_old_hk_acct); + + assert_eq!( + new_null_stake_datura, + old.old_null_stake_datura + .saturating_add(old.old_migration_stake_datura) + ); + } + _ => { + log::warn!("Failed to get account id from ss58 for datura hotkeys"); + return Err("Failed to get account id from ss58 for datura hotkeys".into()); + } + } + + match migration_ck_account { + Ok(migration_ck_acct) => { + // Check the migration key has stake with both *old* hotkeys + assert_eq!( + TotalColdkeyStake::::get(migration_ck_acct), + null_stake_total + ); + } + _ => { + log::warn!("Failed to get account id from ss58 for migration coldkey"); + return Err("Failed to get account id from ss58 for migration coldkey".into()); + } + } + + // Check the total issuance is the SAME following migration (no TAO issued) + let expected_total_issuance = old.total_issuance_before; + let expected_total_stake = old.total_stake_before; + assert_eq!(Self::get_total_issuance(), expected_total_issuance); + + // Check total stake is the SAME following the migration (no new TAO staked) + assert_eq!(TotalStake::::get(), expected_total_stake); + // Check the total stake maps are updated following the migration (removal of old null_account stake entries) + assert_eq!(TotalColdkeyStake::::get(null_account), 0); + + // Check staking hotkeys is updated + assert_eq!(StakingHotkeys::::get(null_account), vec![]); + + Ok(()) + } +} + +pub mod migration { + use frame_support::pallet_prelude::Weight; + use frame_support::traits::OnRuntimeUpgrade; + use sp_core::Get; + + use super::*; + + pub struct Migration(PhantomData); + + #[cfg(feature = "try-runtime")] + fn get_old_storage_values() -> Result { + log::info!("Getting old storage values for migration"); + + let null_account = &DefaultAccount::::get(); + let migration_coldkey = "5GeRjQYsobRWFnrbBmGe5ugme3rfnDVF69N45YtdBpUFsJG8"; + let migration_account = &get_account_id_from_ss58::(migration_coldkey); + + let taostats_old_hotkey = "5Hddm3iBFD2GLT5ik7LZnT3XJUnRnN8PoeCFgGQgawUVKNm8"; + let taostats_new_hotkey = "5GKH9FPPnWSUoeeTJp19wVtd84XqFW4pyK2ijV2GsFbhTrP1"; + + let taostats_old_hk_account = &get_account_id_from_ss58::(taostats_old_hotkey); + let taostats_new_hk_account = &get_account_id_from_ss58::(taostats_new_hotkey); + + let total_issuance_before = crate::Pallet::::get_total_issuance(); + let mut expected_taostats_new_hk_pending_emission: u64 = 0; + let mut expected_datura_new_hk_pending_emission: u64 = 0; + let (old_null_stake_taostats, old_migration_stake_taostats) = match ( + taostats_old_hk_account, + taostats_new_hk_account, + migration_account, + ) { + (Ok(taostats_old_hk_acct), Ok(taostats_new_hk_acct), Ok(migration_acct)) => { + expected_taostats_new_hk_pending_emission = + expected_taostats_new_hk_pending_emission + .saturating_add(PendingdHotkeyEmission::::get(taostats_old_hk_acct)) + .saturating_add(PendingdHotkeyEmission::::get(taostats_new_hk_acct)); + + Ok::<(u64, u64), sp_runtime::TryRuntimeError>(( + crate::Pallet::::get_stake_for_coldkey_and_hotkey( + null_account, + taostats_old_hk_acct, + ), + crate::Pallet::::get_stake_for_coldkey_and_hotkey( + migration_acct, + taostats_old_hk_acct, + ), + )) + } + _ => { + log::warn!("Failed to get account id from ss58 for taostats hotkeys"); + Err("Failed to get account id from ss58 for taostats hotkeys".into()) + } + }?; + + let datura_old_hotkey = "5FKstHjZkh4v3qAMSBa1oJcHCLjxYZ8SNTSz1opTv4hR7gVB"; + let datura_new_hotkey = "5GP7c3fFazW9GXK8Up3qgu2DJBk8inu4aK9TZy3RuoSWVCMi"; + + let datura_old_hk_account = &get_account_id_from_ss58::(datura_old_hotkey); + let datura_new_hk_account = &get_account_id_from_ss58::(datura_new_hotkey); + + let (old_null_stake_datura, old_migration_stake_datura) = match ( + datura_old_hk_account, + datura_new_hk_account, + migration_account, + ) { + (Ok(datura_old_hk_acct), Ok(datura_new_hk_acct), Ok(migration_acct)) => { + expected_datura_new_hk_pending_emission = expected_datura_new_hk_pending_emission + .saturating_add(PendingdHotkeyEmission::::get(datura_old_hk_acct)) + .saturating_add(PendingdHotkeyEmission::::get(datura_new_hk_acct)); + + Ok::<(u64, u64), sp_runtime::TryRuntimeError>(( + crate::Pallet::::get_stake_for_coldkey_and_hotkey( + null_account, + datura_old_hk_acct, + ), + crate::Pallet::::get_stake_for_coldkey_and_hotkey( + migration_acct, + datura_old_hk_acct, + ), + )) + } + _ => { + log::warn!("Failed to get account id from ss58 for datura hotkeys"); + Err("Failed to get account id from ss58 for datura hotkeys".into()) + } + }?; + + let total_stake_before: u64 = crate::Pallet::::get_total_stake(); + + let result = v0::OldStorage { + total_issuance_before, + total_stake_before, + expected_taostats_new_hk_pending_emission, + expected_datura_new_hk_pending_emission, + old_migration_stake_taostats, + old_null_stake_taostats, + old_migration_stake_datura, + old_null_stake_datura, + }; + + log::info!("Got old storage values for migration"); + + Ok(result) + } + + impl OnRuntimeUpgrade for Migration { + /// Runs the migration to fix the pending emissions. + #[cfg(feature = "try-runtime")] + fn pre_upgrade() -> Result, sp_runtime::TryRuntimeError> { + use codec::Encode; + + // Get the old storage values + match get_old_storage_values::() { + Ok(old_storage) => { + log::info!("Successfully got old storage values for migration"); + let encoded = old_storage.encode(); + + Ok(encoded) + } + Err(e) => { + log::error!("Failed to get old storage values for migration: {:?}", e); + Err("Failed to get old storage values for migration".into()) + } + } + } + + // Runs the migrate function for the fix_pending_emission migration + fn on_runtime_upgrade() -> Weight { + let migration_name = b"fix_pending_emission".to_vec(); + + // Initialize the weight with one read operation. + let mut weight = T::DbWeight::get().reads(1); + + // Check if the migration has already run + if HasMigrationRun::::get(&migration_name) { + log::info!( + "Migration '{:?}' has already run. Skipping.", + migration_name + ); + return Weight::zero(); + } + + log::info!( + "Running migration '{}'", + String::from_utf8_lossy(&migration_name) + ); + + // Run the migration + weight.saturating_accrue( + migrations::migrate_fix_pending_emission::do_migrate_fix_pending_emission::(), + ); + + // Mark the migration as completed + HasMigrationRun::::insert(&migration_name, true); + weight.saturating_accrue(T::DbWeight::get().writes(1)); + + log::info!( + "Migration '{:?}' completed. Marked in storage.", + String::from_utf8_lossy(&migration_name) + ); + + // Return the migration weight. + weight + } + + /// Performs post-upgrade checks to ensure the migration was successful. + /// + /// This function is only compiled when the "try-runtime" feature is enabled. + #[cfg(feature = "try-runtime")] + fn post_upgrade(state: Vec) -> Result<(), sp_runtime::TryRuntimeError> { + use codec::Decode; + + let old_storage: v0::OldStorage = + v0::OldStorage::decode(&mut &state[..]).map_err(|_| { + sp_runtime::TryRuntimeError::Other("Failed to decode old value from storage") + })?; + + // Verify that all null stake invariants are satisfied after the migration + crate::Pallet::::check_null_stake_invariants(old_storage)?; + + Ok(()) + } + } +} diff --git a/pallets/subtensor/src/migrations/mod.rs b/pallets/subtensor/src/migrations/mod.rs index a0ee659981..6341ca0c85 100644 --- a/pallets/subtensor/src/migrations/mod.rs +++ b/pallets/subtensor/src/migrations/mod.rs @@ -4,6 +4,7 @@ pub mod migrate_commit_reveal_v2; pub mod migrate_create_root_network; pub mod migrate_delete_subnet_21; pub mod migrate_delete_subnet_3; +pub mod migrate_fix_pending_emission; pub mod migrate_fix_total_coldkey_stake; pub mod migrate_init_total_issuance; pub mod migrate_populate_owned_hotkeys; diff --git a/pallets/subtensor/src/utils/try_state.rs b/pallets/subtensor/src/utils/try_state.rs index 3c01a9b64a..385a21bdd0 100644 --- a/pallets/subtensor/src/utils/try_state.rs +++ b/pallets/subtensor/src/utils/try_state.rs @@ -17,7 +17,7 @@ impl Pallet { // Calculate the total staked amount let mut total_staked: u64 = 0; - for (_account, _netuid, stake) in Stake::::iter() { + for (_hotkey, _coldkey, stake) in Stake::::iter() { total_staked = total_staked.saturating_add(stake); } diff --git a/pallets/subtensor/tests/migration.rs b/pallets/subtensor/tests/migration.rs index 4ddef882c1..1c5076beb7 100644 --- a/pallets/subtensor/tests/migration.rs +++ b/pallets/subtensor/tests/migration.rs @@ -11,9 +11,10 @@ use frame_support::{ use frame_system::Config; use mock::*; use pallet_subtensor::*; -use sp_core::{H256, U256}; +use sp_core::{crypto::Ss58Codec, H256, U256}; use sp_io::hashing::twox_128; use sp_runtime::traits::Zero; +use substrate_fixed::types::extra::U2; #[test] fn test_initialise_ti() { @@ -528,3 +529,158 @@ fn test_migrate_commit_reveal_2() { assert!(!weight.is_zero(), "Migration weight should be non-zero"); }); } + +fn run_pending_emissions_migration_and_check( + migration_name: &'static str, +) -> frame_support::weights::Weight { + use frame_support::traits::OnRuntimeUpgrade; + + // Execute the migration and store its weight + let weight: frame_support::weights::Weight = + pallet_subtensor::migrations::migrate_fix_pending_emission::migration::Migration::::on_runtime_upgrade(); + + // Check if the migration has been marked as completed + assert!(HasMigrationRun::::get( + migration_name.as_bytes().to_vec() + )); + + // Return the weight of the executed migration + weight +} + +fn get_account_id_from_ss58(ss58_str: &str) -> U256 { + let account_id = sp_core::crypto::AccountId32::from_ss58check(ss58_str).unwrap(); + let account_id = AccountId::decode(&mut account_id.as_ref()).unwrap(); + account_id +} + +// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --package pallet-subtensor --test migration -- test_migrate_fix_pending_emissions --exact --nocapture +#[test] +fn test_migrate_fix_pending_emissions() { + new_test_ext(1).execute_with(|| { + let migration_name = "fix_pending_emission"; + + let null_account = &U256::from(0); // The null account + + let taostats_old_hotkey = "5Hddm3iBFD2GLT5ik7LZnT3XJUnRnN8PoeCFgGQgawUVKNm8"; + let taostats_new_hotkey = "5GKH9FPPnWSUoeeTJp19wVtd84XqFW4pyK2ijV2GsFbhTrP1"; + + let taostats_old_hk_account: &AccountId = &get_account_id_from_ss58(taostats_old_hotkey); + let taostats_new_hk_account: &AccountId = &get_account_id_from_ss58(taostats_new_hotkey); + + let datura_old_hotkey = "5FKstHjZkh4v3qAMSBa1oJcHCLjxYZ8SNTSz1opTv4hR7gVB"; + let datura_new_hotkey = "5GP7c3fFazW9GXK8Up3qgu2DJBk8inu4aK9TZy3RuoSWVCMi"; + + let datura_old_hk_account: &AccountId = &get_account_id_from_ss58(datura_old_hotkey); + let datura_new_hk_account: &AccountId = &get_account_id_from_ss58(datura_new_hotkey); + + let migration_coldkey = "5GeRjQYsobRWFnrbBmGe5ugme3rfnDVF69N45YtdBpUFsJG8"; + let migration_account: &AccountId = &get_account_id_from_ss58(migration_coldkey); + + // "Issue" the TAO we're going to insert to stake + let null_stake_datura = 123_456_789; + let null_stake_tao_stats = 123_456_789; + let null_stake_total = null_stake_datura + null_stake_tao_stats; + SubtensorModule::set_total_issuance(null_stake_total); + TotalStake::::put(null_stake_total); + TotalColdkeyStake::::insert(null_account, null_stake_total); + TotalHotkeyStake::::insert(datura_old_hk_account, null_stake_datura); + TotalHotkeyStake::::insert(taostats_old_hk_account, null_stake_tao_stats); + + // Setup the old Datura hotkey with a pending emission + PendingdHotkeyEmission::::insert(datura_old_hk_account, 10_000); + // Setup the NEW Datura hotkey with a pending emission + PendingdHotkeyEmission::::insert(datura_new_hk_account, 123_456_789); + Stake::::insert(datura_old_hk_account, null_account, null_stake_datura); + let expected_datura_new_hk_pending_emission: u64 = 123_456_789 + 10_000; + + // Setup the old TaoStats hotkey with a pending emission + PendingdHotkeyEmission::::insert(taostats_old_hk_account, 987_654); + // Setup the new TaoStats hotkey with a pending emission + PendingdHotkeyEmission::::insert(taostats_new_hk_account, 100_000); + // Setup the old TaoStats hotkey with a null-key stake entry + Stake::::insert(taostats_old_hk_account, null_account, null_stake_tao_stats); + let expected_taostats_new_hk_pending_emission: u64 = 987_654 + 100_000; + + let total_issuance_before = SubtensorModule::get_total_issuance(); + + // Run migration + let first_weight = run_pending_emissions_migration_and_check(migration_name); + assert!(first_weight != Weight::zero()); + + // Check the pending emission is added to new Datura hotkey + assert_eq!( + PendingdHotkeyEmission::::get(datura_new_hk_account), + expected_datura_new_hk_pending_emission + ); + + // Check the pending emission is added to new the TaoStats hotkey + assert_eq!( + PendingdHotkeyEmission::::get(taostats_new_hk_account), + expected_taostats_new_hk_pending_emission + ); + + // Check the pending emission is removed from old ones + assert_eq!( + PendingdHotkeyEmission::::get(datura_old_hk_account), + 0 + ); + + assert_eq!( + PendingdHotkeyEmission::::get(taostats_old_hk_account), + 0 + ); + + // Check the stake entry is removed + assert_eq!(Stake::::get(datura_old_hk_account, null_account), 0); + assert_eq!(Stake::::get(taostats_old_hk_account, null_account), 0); + + // Check the total issuance is the SAME following migration (no TAO issued) + let expected_total_issuance = total_issuance_before; + assert_eq!( + SubtensorModule::get_total_issuance(), + expected_total_issuance + ); + + // Check total stake is the SAME following the migration (no new TAO staked) + assert_eq!(TotalStake::::get(), expected_total_issuance); + // Check the total stake maps are updated following the migration (removal of old null_account stake entries) + assert_eq!(TotalColdkeyStake::::get(null_account), 0); + assert_eq!( + SubtensorModule::get_stake_for_coldkey_and_hotkey(null_account, datura_old_hk_account), + 0 + ); + assert_eq!( + SubtensorModule::get_stake_for_coldkey_and_hotkey( + null_account, + taostats_old_hk_account + ), + 0 + ); + + // Check staking hotkeys is updated + assert_eq!(StakingHotkeys::::get(null_account), vec![]); + + // Check the migration key has stake with both *old* hotkeys + assert_eq!( + SubtensorModule::get_stake_for_coldkey_and_hotkey( + migration_account, + datura_old_hk_account + ), + null_stake_datura + ); + assert_eq!( + SubtensorModule::get_stake_for_coldkey_and_hotkey( + migration_account, + taostats_old_hk_account + ), + null_stake_tao_stats + ); + assert_eq!( + TotalColdkeyStake::::get(migration_account), + null_stake_total + ); + assert!(StakingHotkeys::::get(migration_account).contains(datura_old_hk_account)); + assert!(StakingHotkeys::::get(migration_account).contains(taostats_old_hk_account)); + }) +} diff --git a/pallets/subtensor/tests/swap_hotkey.rs b/pallets/subtensor/tests/swap_hotkey.rs index 206cc324f1..0901e7bdae 100644 --- a/pallets/subtensor/tests/swap_hotkey.rs +++ b/pallets/subtensor/tests/swap_hotkey.rs @@ -9,6 +9,7 @@ use mock::*; use pallet_subtensor::*; use sp_core::H256; use sp_core::U256; +use sp_runtime::SaturatedConversion; // SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_owner --exact --nocapture #[test] @@ -1245,3 +1246,91 @@ fn test_swap_child_hotkey_childkey_maps() { ); }) } + +// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_hotkey_swap_stake_delta --exact --nocapture +#[test] +fn test_hotkey_swap_stake_delta() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(3); + let new_hotkey = U256::from(4); + let coldkey = U256::from(7); + + let coldkeys = [U256::from(1), U256::from(2), U256::from(5)]; + + let mut weight = Weight::zero(); + + // Set up initial state + // Add stake delta for each coldkey and the old_hotkey + for &coldkey in coldkeys.iter() { + StakeDeltaSinceLastEmissionDrain::::insert( + old_hotkey, + coldkey, + (123 + coldkey.saturated_into::()), + ); + + StakingHotkeys::::insert(coldkey, vec![old_hotkey]); + } + + // Add stake delta for one coldkey and the new_hotkey + StakeDeltaSinceLastEmissionDrain::::insert(new_hotkey, coldkeys[0], 456); + // Add corresponding StakingHotkeys + StakingHotkeys::::insert(coldkeys[0], vec![old_hotkey, new_hotkey]); + + // Perform the swap + SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight); + + // Ensure the stake delta is correctly transferred for each coldkey + // -- coldkey[0] maintains its stake delta from the new_hotkey and the old_hotkey + assert_eq!( + StakeDeltaSinceLastEmissionDrain::::get(new_hotkey, coldkeys[0]), + 123 + coldkeys[0].saturated_into::() + 456 + ); + // -- coldkey[1..] maintains its stake delta from the old_hotkey + for &coldkey in coldkeys[1..].iter() { + assert_eq!( + StakeDeltaSinceLastEmissionDrain::::get(new_hotkey, coldkey), + 123 + coldkey.saturated_into::() + ); + assert!(!StakeDeltaSinceLastEmissionDrain::::contains_key( + old_hotkey, coldkey + )); + } + }); +} + +// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_hotkey_with_pending_emissions --exact --nocapture +#[test] +fn test_swap_hotkey_with_pending_emissions() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(1); + let new_hotkey = U256::from(2); + let coldkey = U256::from(3); + let netuid = 0u16; + let mut weight = Weight::zero(); + + let pending_emission = 123_456_789u64; + + // Set up initial state + add_network(netuid, 0, 1); + + // Set up pending emissions + PendingdHotkeyEmission::::insert(old_hotkey, pending_emission); + // Verify the pending emissions are set + assert_eq!( + PendingdHotkeyEmission::::get(old_hotkey), + pending_emission + ); + // Verify the new hotkey does not have any pending emissions + assert!(!PendingdHotkeyEmission::::contains_key(new_hotkey)); + + // Perform the swap + SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight); + + // Verify the pending emissions are transferred + assert_eq!( + PendingdHotkeyEmission::::get(new_hotkey), + pending_emission + ); + assert!(!PendingdHotkeyEmission::::contains_key(old_hotkey)); + }); +}