Merge pull request #614 from opentensor/feat/cold_key_swap

garrett-opentensor · web-flow · commit a62975c8d97c · 2024-07-04T17:30:47.000-07:00
adds tests fixes remove hotkey
diff --git a/pallets/subtensor/src/events.rs b/pallets/subtensor/src/events.rs
@@ -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,
diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs
@@ -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 ------------------------------------------------------------
diff --git a/pallets/subtensor/src/staking.rs b/pallets/subtensor/src/staking.rs
@@ -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
@@ -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);
@@ -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,
         });
 
diff --git a/pallets/subtensor/tests/staking.rs b/pallets/subtensor/tests/staking.rs
@@ -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
         ));
 
@@ -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
@@ -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,
         );
 
@@ -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
         ));
 
@@ -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(),
@@ -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
         ));
 
@@ -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(),
diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs
@@ -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,
