fix migration tests

Author: const

pallets/subtensor/src/lib.rs

@@ -1073,9 +1073,12 @@ pub mod pallet {
         (H256, u64),
         OptionQuery,
     >;
+    
     /// ==================
     /// ==== Genesis =====
     /// ==================
+    #[pallet::storage] // --- Storage for migration run status
+    pub type HasMigrationRun<T: Config> = StorageMap<_, Identity, Vec<u8>, bool, ValueQuery>;
 
     #[pallet::genesis_config]
     pub struct GenesisConfig<T: Config> {

pallets/subtensor/src/macros/dispatches.rs

@@ -640,17 +640,17 @@ mod dispatches {
         }
 
         /// The extrinsic for user to change its hotkey
-        ///#[pallet::call_index(70)]
-        ///#[pallet::weight((Weight::from_parts(1_940_000_000, 0)
-        ///.saturating_add(T::DbWeight::get().reads(272))
-        ///.saturating_add(T::DbWeight::get().writes(527)), DispatchClass::Operational, Pays::No))]
-        ///pub fn swap_hotkey(
-        ///    origin: OriginFor<T>,
-        ///    hotkey: T::AccountId,
-        ///    new_hotkey: T::AccountId,
-        ///) -> DispatchResultWithPostInfo {
-        ///    Self::do_swap_hotkey(origin, &hotkey, &new_hotkey)
-        ///}
+        #[pallet::call_index(70)]
+        #[pallet::weight((Weight::from_parts(1_940_000_000, 0)
+        .saturating_add(T::DbWeight::get().reads(272))
+        .saturating_add(T::DbWeight::get().writes(527)), DispatchClass::Operational, Pays::No))]
+        pub fn swap_hotkey(
+            origin: OriginFor<T>,
+            hotkey: T::AccountId,
+            new_hotkey: T::AccountId,
+        ) -> DispatchResultWithPostInfo {
+            Self::do_swap_hotkey(origin, &hotkey, &new_hotkey)
+        }
 
         /// The extrinsic for user to change the coldkey associated with their account.
         ///

pallets/subtensor/src/migrations/migrate_fix_total_coldkey_stake.rs

@@ -2,9 +2,10 @@ use super::*;
 use frame_support::{
     pallet_prelude::{Identity, OptionQuery},
     storage_alias,
-    traits::{Get, GetStorageVersion, StorageVersion},
+    traits::{Get, StorageVersion},
     weights::Weight,
 };
+use alloc::string::String;
 use sp_std::vec::Vec;
 
 // TODO (camfairchild): TEST MIGRATION
@@ -50,24 +51,41 @@ pub fn do_migrate_fix_total_coldkey_stake<T: Config>() -> Weight {
 }
 // Public migrate function to be called by Lib.rs on upgrade.
 pub fn migrate_fix_total_coldkey_stake<T: Config>() -> Weight {
-    let current_storage_version: u16 = 7;
-    let next_storage_version: u16 = 8;
+    let migration_name = b"fix_total_coldkey_stake_v7".to_vec();
 
     // Initialize the weight with one read operation.
     let mut weight = T::DbWeight::get().reads(1);
 
-    // Grab the current on-chain storage version.
-    // Cant fail on retrieval.
-    let onchain_version = Pallet::<T>::on_chain_storage_version();
-
-    // Only run this migration on storage version 6.
-    if onchain_version == current_storage_version {
-        weight = weight.saturating_add(do_migrate_fix_total_coldkey_stake::<T>());
-        // Cant fail on insert.
-        StorageVersion::new(next_storage_version).put::<Pallet<T>>();
-        weight.saturating_accrue(T::DbWeight::get().writes(1));
+    // Check if the migration has already run
+    if HasMigrationRun::<T>::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 = weight.saturating_add(do_migrate_fix_total_coldkey_stake::<T>());
+
+    // Mark the migration as completed
+    HasMigrationRun::<T>::insert(&migration_name, true);
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    // Set the storage version to 7
+    StorageVersion::new(7).put::<Pallet<T>>();
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    log::info!(
+        "Migration '{:?}' completed. Storage version set to 7.",
+        String::from_utf8_lossy(&migration_name)
+    );
+
     // Return the migration weight.
     weight
 }