add stake balance as equal

Author: const

pallets/subtensor/src/swap.rs

@@ -199,6 +199,16 @@ impl<T: Config> Pallet<T> {
         }
     }
 
+    pub fn meets_min_allowed_coldkey_balance( coldkey: &T::AccountId ) -> bool {
+        let all_staked_keys: Vec<T::AccountId> = StakingHotkeys::<T>::get(coldkey);
+        let mut total_staking_balance: u64 = 0;
+        for hotkey in all_staked_keys {
+            total_staking_balance += Self::get_stake_for_coldkey_and_hotkey(&coldkey, &hotkey);
+        }
+        total_staking_balance += Self::get_coldkey_balance(&coldkey);
+        total_staking_balance >= MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP
+    }
+
     /// Schedules a coldkey swap to a new coldkey with arbitration.
     ///
     /// # Arguments
@@ -238,7 +248,7 @@ impl<T: Config> Pallet<T> {
 
         // Check minimum amount of TAO (1 TAO)
         ensure!(
-            Self::get_coldkey_balance(old_coldkey) >= MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP,
+            Self::meets_min_allowed_coldkey_balance(&old_coldkey),
             Error::<T>::InsufficientBalanceToPerformColdkeySwap
         );
 

pallets/subtensor/tests/staking.rs

@@ -3980,3 +3980,41 @@ fn next_block() {
     System::on_initialize(System::block_number());
     SubtensorModule::on_initialize(System::block_number());
 }
+
+// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test staking -- test_coldkey_meets_enough --exact --nocapture
+#[test]
+fn test_coldkey_meets_enough() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(1);
+        let new_coldkey = U256::from(2);
+        let hotkey = U256::from(2);
+        let netuid = 1u16;
+        add_network(netuid, 13, 0);
+        register_ok_neuron(netuid, hotkey, coldkey, 0);
+        let current_block = SubtensorModule::get_current_block_as_u64();
+        let (work1, nonce1) = generate_valid_pow(&coldkey, current_block, U256::from(10_000_000u64));
+        assert_err!(
+            SubtensorModule::do_schedule_coldkey_swap(
+                &coldkey.clone(),
+                &new_coldkey,
+                work1.to_fixed_bytes().to_vec(),
+                current_block,
+                nonce1
+            ),
+            Error::<Test>::InsufficientBalanceToPerformColdkeySwap
+        );
+        SubtensorModule::add_balance_to_coldkey_account(
+            &coldkey,
+            MIN_BALANCE_TO_PERFORM_COLDKEY_SWAP,
+        );
+        assert_ok!(SubtensorModule::do_schedule_coldkey_swap(
+            &coldkey.clone(),
+            &new_coldkey,
+            work1.to_fixed_bytes().to_vec(),
+            current_block,
+            nonce1
+        ));
+      
+
+    });
+}

pallets/subtensor/tests/swap.rs

@@ -1367,20 +1367,12 @@ fn test_coldkey_has_associated_hotkeys() {
         // Setup initial state
         add_network(netuid, 13, 0);
         register_ok_neuron(netuid, hotkey, coldkey, 0);
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey, 1000);
 
-        // Check if coldkey has associated hotkeys
-        assert!(SubtensorModule::coldkey_has_associated_hotkeys(&coldkey));
-
-        // Check for a coldkey without associated hotkeys
-        let unassociated_coldkey = U256::from(3);
-        assert!(!SubtensorModule::coldkey_has_associated_hotkeys(
-            &unassociated_coldkey
-        ));
     });
 }
 
 
-
 // SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test swap -- test_coldkey_swap_total --exact --nocapture
 #[test]
 fn test_coldkey_swap_total() {