From 7062a3e5263128c3f8ee65e391427db0409d33d2 Mon Sep 17 00:00:00 2001 From: Samuel Dare Date: Thu, 4 Jul 2024 23:31:49 +0400 Subject: [PATCH] chore: additional tests --- pallets/subtensor/src/swap.rs | 1 + pallets/subtensor/tests/swap.rs | 95 ++++++++++++++++++++------------- 2 files changed, 59 insertions(+), 37 deletions(-) diff --git a/pallets/subtensor/src/swap.rs b/pallets/subtensor/src/swap.rs index b6aed8b72..49af72a3f 100644 --- a/pallets/subtensor/src/swap.rs +++ b/pallets/subtensor/src/swap.rs @@ -128,6 +128,7 @@ impl Pallet { let block: u64 = Self::get_current_block_as_u64(); // Swap coldkey references in storage maps + // NOTE The order of these calls is important Self::swap_total_coldkey_stake(old_coldkey, new_coldkey, &mut weight); Self::swap_stake_for_coldkey(old_coldkey, new_coldkey, &mut weight); Self::swap_owner_for_coldkey(old_coldkey, new_coldkey, &mut weight); diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index 679393b6d..324ca3f10 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1054,44 +1054,51 @@ fn test_do_swap_coldkey_success() { new_test_ext(1).execute_with(|| { let old_coldkey = U256::from(1); let new_coldkey = U256::from(2); - let hotkey = U256::from(3); + let hotkey1 = U256::from(3); + let hotkey2 = U256::from(4); let netuid = 1u16; - let stake_amount = 1000u64; - let free_balance = 12345; + let stake_amount1 = 1000u64; + let stake_amount2 = 2000u64; + let free_balance_old = 12345u64; // Setup initial state add_network(netuid, 13, 0); - register_ok_neuron(netuid, hotkey, old_coldkey, 0); - SubtensorModule::add_balance_to_coldkey_account(&old_coldkey, stake_amount + free_balance); + register_ok_neuron(netuid, hotkey1, old_coldkey, 0); + register_ok_neuron(netuid, hotkey2, old_coldkey, 0); - // Add stake to the neuron + // Add balance to old coldkey + SubtensorModule::add_balance_to_coldkey_account( + &old_coldkey, + stake_amount1 + stake_amount2 + free_balance_old, + ); + + // Add stake to the neurons + assert_ok!(SubtensorModule::add_stake( + <::RuntimeOrigin>::signed(old_coldkey), + hotkey1, + stake_amount1 + )); assert_ok!(SubtensorModule::add_stake( <::RuntimeOrigin>::signed(old_coldkey), - hotkey, - stake_amount + hotkey2, + stake_amount2 )); - log::info!( - "TotalColdkeyStake::::get(old_coldkey): {:?}", - TotalColdkeyStake::::get(old_coldkey) + // Verify initial stakes and balances + assert_eq!( + TotalColdkeyStake::::get(old_coldkey), + stake_amount1 + stake_amount2 ); - log::info!( - "Stake::::get(old_coldkey, hotkey): {:?}", - Stake::::get(hotkey, old_coldkey) + assert_eq!(Stake::::get(hotkey1, old_coldkey), stake_amount1); + assert_eq!(Stake::::get(hotkey2, old_coldkey), stake_amount2); + assert_eq!( + OwnedHotkeys::::get(old_coldkey), + vec![hotkey1, hotkey2] + ); + assert_eq!( + SubtensorModule::get_coldkey_balance(&old_coldkey), + free_balance_old ); - - // Verify initial stake - assert_eq!(TotalColdkeyStake::::get(old_coldkey), stake_amount); - assert_eq!(Stake::::get(hotkey, old_coldkey), stake_amount); - - assert_eq!(OwnedHotkeys::::get(old_coldkey), vec![hotkey]); - assert!(!OwnedHotkeys::::get(new_coldkey).contains(hotkey)); - - // Get coldkey free balance before swap - let balance = SubtensorModule::get_coldkey_balance(&old_coldkey); - assert_eq!(balance, free_balance); - -let balance_new_coldkey = SubtensorModule::get_coldkey_balance(&new_coldkey); // Perform the swap assert_ok!(SubtensorModule::do_swap_coldkey( @@ -1101,16 +1108,30 @@ let balance_new_coldkey = SubtensorModule::get_coldkey_balance(&new_coldkey); )); // Verify the swap - assert_eq!(Owner::::get(hotkey), new_coldkey); - assert_eq!(TotalColdkeyStake::::get(new_coldkey), stake_amount); + assert_eq!(Owner::::get(hotkey1), new_coldkey); + assert_eq!(Owner::::get(hotkey2), new_coldkey); + assert_eq!( + TotalColdkeyStake::::get(new_coldkey), + stake_amount1 + stake_amount2 + ); assert!(!TotalColdkeyStake::::contains_key(old_coldkey)); - assert_eq!(Stake::::get(hotkey, new_coldkey), stake_amount); - assert!(!Stake::::contains_key(hotkey, old_coldkey)); - assert_eq!(OwnedHotkeys::::get(new_coldkey), vec![hotkey]); + assert_eq!(Stake::::get(hotkey1, new_coldkey), stake_amount1); + assert_eq!(Stake::::get(hotkey2, new_coldkey), stake_amount2); + assert!(!Stake::::contains_key(hotkey1, old_coldkey)); + assert!(!Stake::::contains_key(hotkey2, old_coldkey)); + + // Verify OwnedHotkeys + let new_owned_hotkeys = OwnedHotkeys::::get(new_coldkey); + assert!(new_owned_hotkeys.contains(&hotkey1)); + assert!(new_owned_hotkeys.contains(&hotkey2)); + assert_eq!(new_owned_hotkeys.len(), 2); assert!(!OwnedHotkeys::::contains_key(old_coldkey)); // Verify balance transfer - assert_eq!(SubtensorModule::get_coldkey_balance(&new_coldkey), balance + balance_new_coldkey); + assert_eq!( + SubtensorModule::get_coldkey_balance(&new_coldkey), + free_balance_old + ); assert_eq!(SubtensorModule::get_coldkey_balance(&old_coldkey), 0); // Verify event emission @@ -1174,10 +1195,10 @@ fn test_swap_stake_for_coldkey() { assert_eq!(Stake::::get(hotkey2, new_coldkey), stake_amount2); assert!(!Stake::::contains_key(hotkey1, old_coldkey)); assert!(!Stake::::contains_key(hotkey2, old_coldkey)); -assert_eq!(TotalHotkeyStake::::get(hotkey1), stake_amount1); -assert_eq!(TotalHotkeyStake::::get(hotkey2), stake_amount2); -assert_eq!(TotalStake::::get(), stake_amount1 + stake_amount2); -assert_eq!(TotalIssuance::::get(), stake_amount1 + stake_amount2); + assert_eq!(TotalHotkeyStake::::get(hotkey1), stake_amount1); + assert_eq!(TotalHotkeyStake::::get(hotkey2), stake_amount2); + assert_eq!(TotalStake::::get(), stake_amount1 + stake_amount2); + assert_eq!(TotalIssuance::::get(), stake_amount1 + stake_amount2); // Verify weight update let expected_weight = ::DbWeight::get().reads_writes(3, 4); assert_eq!(weight, expected_weight);