chore: fix tests

Author: Samuel Dare

pallets/subtensor/src/swap.rs

@@ -60,6 +60,16 @@ impl<T: Config> Pallet<T> {
             T::DbWeight::get().reads((TotalNetworks::<T>::get().saturating_add(1u16)) as u64),
         );
 
+        let swap_cost = Self::get_hotkey_swap_cost();
+        log::debug!("Swap cost: {:?}", swap_cost);
+
+        ensure!(
+            Self::can_remove_balance_from_coldkey_account(&coldkey, swap_cost),
+            Error::<T>::NotEnoughBalanceToPaySwapHotKey
+        );
+        let actual_burn_amount = Self::remove_balance_from_coldkey_account(&coldkey, swap_cost)?;
+        Self::burn_tokens(actual_burn_amount);
+
         Self::swap_owner(old_hotkey, new_hotkey, &coldkey, &mut weight);
         Self::swap_total_hotkey_stake(old_hotkey, new_hotkey, &mut weight);
         Self::swap_delegates(old_hotkey, new_hotkey, &mut weight);
@@ -76,7 +86,6 @@ impl<T: Config> Pallet<T> {
         Self::swap_prometheus(old_hotkey, new_hotkey, &netuid_is_member, &mut weight);
 
         Self::swap_total_hotkey_coldkey_stakes_this_interval(old_hotkey, new_hotkey, &mut weight);
-        Self::swap_senate_member(old_hotkey, new_hotkey, &mut weight)?;
 
         Self::set_last_tx_block(&coldkey, block);
         weight.saturating_accrue(T::DbWeight::get().writes(1));
@@ -950,15 +959,15 @@ impl<T: Config> Pallet<T> {
         weight.saturating_accrue(T::DbWeight::get().reads(TotalNetworks::<T>::get() as u64));
     }
 
-    /// Swaps the Senate membership from the old hotkey to the new hotkey if applicable.
     pub fn swap_senate_member(
         old_hotkey: &T::AccountId,
         new_hotkey: &T::AccountId,
         weight: &mut Weight,
     ) -> DispatchResult {
+        weight.saturating_accrue(T::DbWeight::get().reads(1));
         if T::SenateMembers::is_member(old_hotkey) {
             T::SenateMembers::swap_member(old_hotkey, new_hotkey).map_err(|e| e.error)?;
-            weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
+            weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
         }
         Ok(())
     }

pallets/subtensor/tests/swap.rs

@@ -3,7 +3,7 @@
 use codec::Encode;
 use frame_support::weights::Weight;
 use frame_support::{assert_err, assert_noop, assert_ok};
-use frame_system::Config;
+use frame_system::{Config, RawOrigin};
 mod mock;
 use mock::*;
 use pallet_subtensor::*;
@@ -1675,11 +1675,9 @@ fn test_swap_senate_member() {
         let mut weight = Weight::zero();
 
         // Setup: Add old_hotkey as a Senate member
-        assert_ok!(Senate::set_members(
-            RuntimeOrigin::root(),
-            vec![old_hotkey],
-            None,
-            0
+        assert_ok!(SenateMembers::add_member(
+            RawOrigin::Root.into(),
+            old_hotkey
         ));
 
         // Test 1: Successful swap
@@ -1710,29 +1708,5 @@ fn test_swap_senate_member() {
         // Verify weight update (should only have read operations)
         let expected_weight = <Test as frame_system::Config>::DbWeight::get().reads(1);
         assert_eq!(weight, expected_weight);
-
-        // Reset weight for next test
-        weight = Weight::zero();
-
-        // Test 3: Setup for swap to an existing member
-        let another_member = U256::from(4);
-        assert_ok!(Senate::set_members(
-            RuntimeOrigin::root(),
-            vec![new_hotkey, another_member],
-            None,
-            1
-        ));
-
-        // Test 4: Setup for swap with maximum members reached
-        let mut members = vec![new_hotkey, another_member];
-        for i in 5..=SenateMaxMembers::get() {
-            members.push(U256::from(i as u64));
-        }
-        assert_ok!(Senate::set_members(
-            RuntimeOrigin::root(),
-            members.clone(),
-            None,
-            2
-        ));
     });
 }