Merge pull request #645 from opentensor/fix/total-coldkey-stake-migration-clear-first

Fix/total coldkey stake migration; clear map first

Author: distributedstatemachine

pallets/subtensor/src/migration.rs

@@ -37,6 +37,13 @@ pub fn do_migrate_fix_total_coldkey_stake<T: Config>() -> 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::<T>::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::<T>::iter() {
         // Init the zero value.

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::<Test>::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::<Test>::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 =

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::<Test>::insert(hotkey0, coldkey);
+        Owner::<Test>::insert(hotkey1, coldkey);
+        OwnedHotkeys::<Test>::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::<Test>::get(coldkey).len(), 1);
+        assert!(StakingHotkeys::<Test>::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::<Test>::get(coldkey).len(), 2);
+        let final_map = StakingHotkeys::<Test>::get(coldkey);
+        assert!(final_map.contains(&hotkey0));
+        assert!(final_map.contains(&hotkey1));
+    })
+}