diff --git a/pallets/subtensor/src/migration.rs b/pallets/subtensor/src/migration.rs index 20e91237fc..866ff08fd1 100644 --- a/pallets/subtensor/src/migration.rs +++ b/pallets/subtensor/src/migration.rs @@ -37,6 +37,13 @@ pub fn do_migrate_fix_total_coldkey_stake() -> Weight { // Initialize the weight with one read operation. let mut weight = T::DbWeight::get().reads(1); + // Clear everything from the map first, no limit (u32::MAX) + let removal_results = TotalColdkeyStake::::clear(u32::MAX, None); + // 1 read/write per removal + let entries_removed: u64 = removal_results.backend.into(); + weight = + weight.saturating_add(T::DbWeight::get().reads_writes(entries_removed, entries_removed)); + // Iterate through all staking hotkeys. for (coldkey, hotkey_vec) in StakingHotkeys::::iter() { // Init the zero value. diff --git a/pallets/subtensor/tests/migration.rs b/pallets/subtensor/tests/migration.rs index 8323cab8a2..d47155862f 100644 --- a/pallets/subtensor/tests/migration.rs +++ b/pallets/subtensor/tests/migration.rs @@ -392,6 +392,23 @@ fn test_migrate_fix_total_coldkey_stake_runs_once() { }) } +// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test migration -- test_migrate_fix_total_coldkey_stake_starts_with_value_no_stake_map_entries --exact --nocapture +#[test] +fn test_migrate_fix_total_coldkey_stake_starts_with_value_no_stake_map_entries() { + new_test_ext(1).execute_with(|| { + let migration_name = "fix_total_coldkey_stake_v7"; + let coldkey = U256::from(0); + TotalColdkeyStake::::insert(coldkey, 123_456_789); + + // Notably, coldkey has no stake map or staking_hotkeys map entries + + let weight = run_migration_and_check(migration_name); + assert!(weight != Weight::zero()); + // Therefore 0 + assert_eq!(TotalColdkeyStake::::get(coldkey), 0); + }) +} + fn run_migration_and_check(migration_name: &'static str) -> frame_support::weights::Weight { // Execute the migration and store its weight let weight: frame_support::weights::Weight = diff --git a/pallets/subtensor/tests/staking.rs b/pallets/subtensor/tests/staking.rs index 5db439e5b7..cb2cac4efe 100644 --- a/pallets/subtensor/tests/staking.rs +++ b/pallets/subtensor/tests/staking.rs @@ -4719,3 +4719,35 @@ fn test_do_schedule_coldkey_swap_regular_user_passes_min_balance() { ); }); } + +// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test staking -- test_emission_creates_staking_hotkeys_entry --exact --nocapture +#[test] +fn test_emission_creates_staking_hotkeys_entry() { + new_test_ext(1).execute_with(|| { + let hotkey0 = U256::from(1); + let hotkey1 = U256::from(2); + + let coldkey = U256::from(3); + + // Add to Owner map + Owner::::insert(hotkey0, coldkey); + Owner::::insert(hotkey1, coldkey); + OwnedHotkeys::::insert(coldkey, vec![hotkey0, hotkey1]); + + // Emit through hotkey + SubtensorModule::emit_inflation_through_hotkey_account(&hotkey0, 0, 1_000); + + // Verify StakingHotkeys has an entry + assert_eq!(StakingHotkeys::::get(coldkey).len(), 1); + assert!(StakingHotkeys::::get(coldkey).contains(&hotkey0)); + + // Try again with another emission on hotkey1 + SubtensorModule::emit_inflation_through_hotkey_account(&hotkey1, 0, 2_000); + + // Verify both hotkeys are now in the map + assert_eq!(StakingHotkeys::::get(coldkey).len(), 2); + let final_map = StakingHotkeys::::get(coldkey); + assert!(final_map.contains(&hotkey0)); + assert!(final_map.contains(&hotkey1)); + }) +}