Merge branch 'feat/rao-2025-01' of github.com:opentensor/subtensor into feat/rao-2025-01

Author: JohnReedV

pallets/subtensor/src/tests/swap_coldkey.rs

@@ -454,16 +454,44 @@ fn test_swap_concurrent_modifications() {
         let new_coldkey = U256::from(2);
         let hotkey = U256::from(3);
         let netuid: u16 = 1;
-        let initial_stake = 100;
-        let additional_stake = 50;
+        let initial_stake = 1_000_000_000_000;
+        let additional_stake = 500_000_000_000;
 
-        StakingHotkeys::<Test>::insert(old_coldkey, vec![hotkey]);
-        Stake::<Test>::insert(hotkey, old_coldkey, initial_stake);
-
-        // Simulate concurrent stake addition
+        // Setup initial state
         add_network(netuid, 1, 1);
-        SubtensorModule::add_balance_to_coldkey_account(&new_coldkey, additional_stake);
+        SubtensorModule::add_balance_to_coldkey_account(
+            &new_coldkey,
+            initial_stake + additional_stake + 1000_000,
+        );
         register_ok_neuron(netuid, hotkey, new_coldkey, 1001000);
+        assert_ok!(SubtensorModule::add_stake(
+            <<Test as Config>::RuntimeOrigin>::signed(new_coldkey),
+            hotkey,
+            netuid,
+            initial_stake
+        ));
+
+        // Verify initial stake
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &hotkey,
+                &new_coldkey,
+                netuid
+            ),
+            initial_stake
+        );
+
+        // Wait some blocks
+        step_block(10);
+
+        // Get stake before swap
+        let stake_before_swap = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey,
+            &new_coldkey,
+            netuid,
+        );
+
+        // Simulate concurrent stake addition
         assert_ok!(SubtensorModule::add_stake(
             <<Test as Config>::RuntimeOrigin>::signed(new_coldkey),
             hotkey,
@@ -478,13 +506,16 @@ fn test_swap_concurrent_modifications() {
             &mut weight
         ));
 
-        assert_eq!(
-            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+        let eps = 500; // RAO
+        assert!(
+            (SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
                 &hotkey,
                 &new_coldkey,
                 netuid
-            ),
-            initial_stake + additional_stake - 1
+            ) as i64
+                - (stake_before_swap + additional_stake) as i64)
+                .abs()
+                <= eps
         );
         assert!(!Alpha::<Test>::contains_key((hotkey, old_coldkey, netuid)));
     });
@@ -843,43 +874,79 @@ fn test_swap_stake_for_coldkey() {
     });
 }
 
-// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test swap_coldkey -- test_swap_staking_hotkeys_for_coldkey --exact --nocapture
+// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --package pallet-subtensor --lib -- tests::swap_coldkey::test_swap_staking_hotkeys_for_coldkey --exact --show-output
 #[test]
 fn test_swap_staking_hotkeys_for_coldkey() {
     new_test_ext(1).execute_with(|| {
-        assert!(false);
+        let old_coldkey = U256::from(1);
+        let new_coldkey = U256::from(2);
+        let other_coldkey = U256::from(3);
+        let hotkey1 = U256::from(4);
+        let hotkey2 = U256::from(5);
+        let stake_amount1 = 1000u64;
+        let stake_amount2 = 2000u64;
+        let total_stake = stake_amount1 + stake_amount2;
+        let mut weight = Weight::zero();
 
-        // let old_coldkey = U256::from(1);
-        // let new_coldkey = U256::from(2);
-        // let hotkey1 = U256::from(3);
-        // let hotkey2 = U256::from(4);
-        // let stake_amount1 = 1000u64;
-        // let stake_amount2 = 2000u64;
-        // let total_stake = stake_amount1 + stake_amount2;
-        // let mut weight = Weight::zero();
+        // Setup initial state
+        // Add a network
+        let netuid = 1u16;
+        add_network(netuid, 1, 0);
+        // Give some balance to old coldkey
+        SubtensorModule::add_balance_to_coldkey_account(
+            &old_coldkey,
+            stake_amount1 + stake_amount2 + 1000_000,
+        );
+        // Register hotkeys
+        register_ok_neuron(netuid, hotkey1, old_coldkey, 0);
+        register_ok_neuron(netuid, hotkey2, other_coldkey, 0);
+        // Make hotkey2 a delegate
+        assert_ok!(SubtensorModule::become_delegate(
+            <<Test as Config>::RuntimeOrigin>::signed(other_coldkey),
+            hotkey2
+        ));
 
-        // // Setup initial state
-        // OwnedHotkeys::<Test>::insert(old_coldkey, vec![hotkey1, hotkey2]);
-        // Stake::<Test>::insert(hotkey1, old_coldkey, stake_amount1);
-        // Stake::<Test>::insert(hotkey2, old_coldkey, stake_amount2);
-        // StakingHotkeys::<Test>::insert(old_coldkey, vec![hotkey1, hotkey2]);
-        // TotalHotkeyStake::<Test>::insert(hotkey1, stake_amount1);
-        // TotalHotkeyStake::<Test>::insert(hotkey2, stake_amount2);
-        // TotalColdkeyStake::<Test>::insert(old_coldkey, total_stake);
+        // Stake to hotkeys
+        assert_ok!(SubtensorModule::add_stake(
+            <<Test as Config>::RuntimeOrigin>::signed(old_coldkey),
+            hotkey1,
+            netuid,
+            stake_amount1
+        ));
+        assert_ok!(SubtensorModule::add_stake(
+            <<Test as Config>::RuntimeOrigin>::signed(old_coldkey),
+            hotkey2,
+            netuid,
+            stake_amount2
+        ));
 
-        // // Set up total issuance
-        // TotalIssuance::<Test>::put(total_stake);
-        // TotalStake::<Test>::put(total_stake);
+        // Verify stakes
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &hotkey1,
+                &old_coldkey,
+                netuid
+            ),
+            stake_amount1
+        );
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &hotkey2,
+                &old_coldkey,
+                netuid
+            ),
+            stake_amount2
+        );
 
-        // // Perform the swap
-        // SubtensorModule::perform_swap_coldkey(&old_coldkey, &new_coldkey, &mut weight);
+        // Perform the swap
+        SubtensorModule::perform_swap_coldkey(&old_coldkey, &new_coldkey, &mut weight);
 
-        // // Verify StakingHotkeys transfer
-        // assert_eq!(
-        //     StakingHotkeys::<Test>::get(new_coldkey),
-        //     vec![hotkey1, hotkey2]
-        // );
-        // assert_eq!(StakingHotkeys::<Test>::get(old_coldkey), vec![]);
+        // Verify StakingHotkeys transfer
+        assert_eq!(
+            StakingHotkeys::<Test>::get(new_coldkey),
+            vec![hotkey1, hotkey2]
+        );
+        assert_eq!(StakingHotkeys::<Test>::get(old_coldkey), vec![]);
     });
 }