From 53288a0449d9ff488e2641f4629220bc449da7c3 Mon Sep 17 00:00:00 2001 From: unconst Date: Wed, 24 Jul 2024 14:07:32 -0500 Subject: [PATCH 1/2] add tests and add swap --- pallets/subtensor/src/swap/swap_hotkey.rs | 34 ++++++ pallets/subtensor/tests/swap_hotkey.rs | 127 ++++++++++++++++++++++ 2 files changed, 161 insertions(+) diff --git a/pallets/subtensor/src/swap/swap_hotkey.rs b/pallets/subtensor/src/swap/swap_hotkey.rs index fb3c33e4d8..faadf5de56 100644 --- a/pallets/subtensor/src/swap/swap_hotkey.rs +++ b/pallets/subtensor/src/swap/swap_hotkey.rs @@ -307,6 +307,40 @@ impl Pallet { weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1)); } + // 11. Swap ChildKeys. + // ChildKeys( parent, netuid ) --> Vec<(proportion,child)> -- the child keys of the parent. + for netuid in Self::get_all_subnet_netuids() { + // Get the children of the old hotkey for this subnet + let my_children: Vec<(u64, T::AccountId)> = ChildKeys::::get(old_hotkey, netuid); + // Remove the old hotkey's child entries + ChildKeys::::remove(old_hotkey, netuid); + // Insert the same child entries for the new hotkey + ChildKeys::::insert(new_hotkey, netuid, my_children); + } + + // 12. Swap ParentKeys. + // ParentKeys( child, netuid ) --> Vec<(proportion,parent)> -- the parent keys of the child. + for netuid in Self::get_all_subnet_netuids() { + // Get the parents of the old hotkey for this subnet + let parents: Vec<(u64, T::AccountId)> = ParentKeys::::get(old_hotkey, netuid); + // Remove the old hotkey's parent entries + ParentKeys::::remove(old_hotkey, netuid); + // Insert the same parent entries for the new hotkey + ParentKeys::::insert(new_hotkey, netuid, parents.clone()); + for (_, parent_key_i) in parents { + // For each parent, update their children list + let mut parent_children: Vec<(u64, T::AccountId)> = ChildKeys::::get(parent_key_i.clone(), netuid); + for child in parent_children.iter_mut() { + // If the child is the old hotkey, replace it with the new hotkey + if child.1 == *old_hotkey { + child.1 = new_hotkey.clone(); + } + } + // Update the parent's children list + ChildKeys::::insert(parent_key_i, netuid, parent_children); + } + } + // Return successful after swapping all the relevant terms. Ok(()) } diff --git a/pallets/subtensor/tests/swap_hotkey.rs b/pallets/subtensor/tests/swap_hotkey.rs index c6a05f2b6f..5339c00088 100644 --- a/pallets/subtensor/tests/swap_hotkey.rs +++ b/pallets/subtensor/tests/swap_hotkey.rs @@ -959,3 +959,130 @@ fn test_swap_hotkey_error_cases() { assert_eq!(Balances::free_balance(coldkey), initial_balance - swap_cost); }); } + + +// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_child_keys --exact --nocapture +#[test] +fn test_swap_child_keys() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(1); + let new_hotkey = U256::from(2); + let coldkey = U256::from(3); + let netuid = 0u16; + let children = vec![(100u64, U256::from(4)), (200u64, U256::from(5))]; + let mut weight = Weight::zero(); + + // Initialize ChildKeys for old_hotkey + add_network(netuid, 1, 0); + ChildKeys::::insert(old_hotkey, netuid, children.clone()); + + // Perform the swap + SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight); + + // Verify the swap + assert_eq!(ChildKeys::::get(new_hotkey, netuid), children); + assert!(ChildKeys::::get(old_hotkey, netuid).is_empty()); + }); +} + +// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_parent_keys --exact --nocapture +#[test] +fn test_swap_parent_keys() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(1); + let new_hotkey = U256::from(2); + let coldkey = U256::from(3); + let netuid = 0u16; + let parents = vec![(100u64, U256::from(4)), (200u64, U256::from(5))]; + let mut weight = Weight::zero(); + + // Initialize ParentKeys for old_hotkey + add_network(netuid, 1, 0); + ParentKeys::::insert(old_hotkey, netuid, parents.clone()); + + // Initialize ChildKeys for parent + ChildKeys::::insert(U256::from(4), netuid, vec![(100u64, old_hotkey)]); + ChildKeys::::insert(U256::from(5), netuid, vec![(200u64, old_hotkey)]); + + // Perform the swap + SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight); + + // Verify ParentKeys swap + assert_eq!(ParentKeys::::get(new_hotkey, netuid), parents); + assert!(ParentKeys::::get(old_hotkey, netuid).is_empty()); + + // Verify ChildKeys update for parents + assert_eq!(ChildKeys::::get(U256::from(4), netuid), vec![(100u64, new_hotkey)]); + assert_eq!(ChildKeys::::get(U256::from(5), netuid), vec![(200u64, new_hotkey)]); + }); +} + +// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_multiple_subnets --exact --nocapture +#[test] +fn test_swap_multiple_subnets() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(1); + let new_hotkey = U256::from(2); + let coldkey = U256::from(3); + let netuid1 = 0u16; + let netuid2 = 1u16; + let children1 = vec![(100u64, U256::from(4)), (200u64, U256::from(5))]; + let children2 = vec![(300u64, U256::from(6))]; + let mut weight = Weight::zero(); + + add_network(netuid1, 1, 0); + add_network(netuid2, 1, 0); + + // Initialize ChildKeys for old_hotkey in multiple subnets + ChildKeys::::insert(old_hotkey, netuid1, children1.clone()); + ChildKeys::::insert(old_hotkey, netuid2, children2.clone()); + + // Perform the swap + SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight); + + // Verify the swap for both subnets + assert_eq!(ChildKeys::::get(new_hotkey, netuid1), children1); + assert_eq!(ChildKeys::::get(new_hotkey, netuid2), children2); + assert!(ChildKeys::::get(old_hotkey, netuid1).is_empty()); + assert!(ChildKeys::::get(old_hotkey, netuid2).is_empty()); + }); +} + +// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_complex_parent_child_structure --exact --nocapture +#[test] +fn test_swap_complex_parent_child_structure() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(1); + let new_hotkey = U256::from(2); + let coldkey = U256::from(3); + let netuid = 0u16; + let parent1 = U256::from(4); + let parent2 = U256::from(5); + let child1 = U256::from(6); + let child2 = U256::from(7); + let mut weight = Weight::zero(); + + add_network(netuid, 1, 0); + + // Set up complex parent-child structure + ParentKeys::::insert(old_hotkey, netuid, vec![(100u64, parent1), (200u64, parent2)]); + ChildKeys::::insert(old_hotkey, netuid, vec![(300u64, child1), (400u64, child2)]); + ChildKeys::::insert(parent1, netuid, vec![(100u64, old_hotkey), (500u64, U256::from(8))]); + ChildKeys::::insert(parent2, netuid, vec![(200u64, old_hotkey), (600u64, U256::from(9))]); + + // Perform the swap + SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight); + + // Verify ParentKeys swap + assert_eq!(ParentKeys::::get(new_hotkey, netuid), vec![(100u64, parent1), (200u64, parent2)]); + assert!(ParentKeys::::get(old_hotkey, netuid).is_empty()); + + // Verify ChildKeys swap + assert_eq!(ChildKeys::::get(new_hotkey, netuid), vec![(300u64, child1), (400u64, child2)]); + assert!(ChildKeys::::get(old_hotkey, netuid).is_empty()); + + // Verify parent's ChildKeys update + assert_eq!(ChildKeys::::get(parent1, netuid), vec![(100u64, new_hotkey), (500u64, U256::from(8))]); + assert_eq!(ChildKeys::::get(parent2, netuid), vec![(200u64, new_hotkey), (600u64, U256::from(9))]); + }); +} \ No newline at end of file From 7107ae08dafa531eea905b97b920a2999ff996e7 Mon Sep 17 00:00:00 2001 From: unconst Date: Wed, 24 Jul 2024 14:07:44 -0500 Subject: [PATCH 2/2] fmt --- pallets/subtensor/src/swap/swap_hotkey.rs | 15 +++---- pallets/subtensor/tests/swap_hotkey.rs | 51 ++++++++++++++++++----- 2 files changed, 48 insertions(+), 18 deletions(-) diff --git a/pallets/subtensor/src/swap/swap_hotkey.rs b/pallets/subtensor/src/swap/swap_hotkey.rs index faadf5de56..366494d117 100644 --- a/pallets/subtensor/src/swap/swap_hotkey.rs +++ b/pallets/subtensor/src/swap/swap_hotkey.rs @@ -310,12 +310,12 @@ impl Pallet { // 11. Swap ChildKeys. // ChildKeys( parent, netuid ) --> Vec<(proportion,child)> -- the child keys of the parent. for netuid in Self::get_all_subnet_netuids() { - // Get the children of the old hotkey for this subnet - let my_children: Vec<(u64, T::AccountId)> = ChildKeys::::get(old_hotkey, netuid); - // Remove the old hotkey's child entries - ChildKeys::::remove(old_hotkey, netuid); - // Insert the same child entries for the new hotkey - ChildKeys::::insert(new_hotkey, netuid, my_children); + // Get the children of the old hotkey for this subnet + let my_children: Vec<(u64, T::AccountId)> = ChildKeys::::get(old_hotkey, netuid); + // Remove the old hotkey's child entries + ChildKeys::::remove(old_hotkey, netuid); + // Insert the same child entries for the new hotkey + ChildKeys::::insert(new_hotkey, netuid, my_children); } // 12. Swap ParentKeys. @@ -329,7 +329,8 @@ impl Pallet { ParentKeys::::insert(new_hotkey, netuid, parents.clone()); for (_, parent_key_i) in parents { // For each parent, update their children list - let mut parent_children: Vec<(u64, T::AccountId)> = ChildKeys::::get(parent_key_i.clone(), netuid); + let mut parent_children: Vec<(u64, T::AccountId)> = + ChildKeys::::get(parent_key_i.clone(), netuid); for child in parent_children.iter_mut() { // If the child is the old hotkey, replace it with the new hotkey if child.1 == *old_hotkey { diff --git a/pallets/subtensor/tests/swap_hotkey.rs b/pallets/subtensor/tests/swap_hotkey.rs index 5339c00088..bff738b86d 100644 --- a/pallets/subtensor/tests/swap_hotkey.rs +++ b/pallets/subtensor/tests/swap_hotkey.rs @@ -960,7 +960,6 @@ fn test_swap_hotkey_error_cases() { }); } - // SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_child_keys --exact --nocapture #[test] fn test_swap_child_keys() { @@ -1012,8 +1011,14 @@ fn test_swap_parent_keys() { assert!(ParentKeys::::get(old_hotkey, netuid).is_empty()); // Verify ChildKeys update for parents - assert_eq!(ChildKeys::::get(U256::from(4), netuid), vec![(100u64, new_hotkey)]); - assert_eq!(ChildKeys::::get(U256::from(5), netuid), vec![(200u64, new_hotkey)]); + assert_eq!( + ChildKeys::::get(U256::from(4), netuid), + vec![(100u64, new_hotkey)] + ); + assert_eq!( + ChildKeys::::get(U256::from(5), netuid), + vec![(200u64, new_hotkey)] + ); }); } @@ -1065,24 +1070,48 @@ fn test_swap_complex_parent_child_structure() { add_network(netuid, 1, 0); // Set up complex parent-child structure - ParentKeys::::insert(old_hotkey, netuid, vec![(100u64, parent1), (200u64, parent2)]); + ParentKeys::::insert( + old_hotkey, + netuid, + vec![(100u64, parent1), (200u64, parent2)], + ); ChildKeys::::insert(old_hotkey, netuid, vec![(300u64, child1), (400u64, child2)]); - ChildKeys::::insert(parent1, netuid, vec![(100u64, old_hotkey), (500u64, U256::from(8))]); - ChildKeys::::insert(parent2, netuid, vec![(200u64, old_hotkey), (600u64, U256::from(9))]); + ChildKeys::::insert( + parent1, + netuid, + vec![(100u64, old_hotkey), (500u64, U256::from(8))], + ); + ChildKeys::::insert( + parent2, + netuid, + vec![(200u64, old_hotkey), (600u64, U256::from(9))], + ); // Perform the swap SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight); // Verify ParentKeys swap - assert_eq!(ParentKeys::::get(new_hotkey, netuid), vec![(100u64, parent1), (200u64, parent2)]); + assert_eq!( + ParentKeys::::get(new_hotkey, netuid), + vec![(100u64, parent1), (200u64, parent2)] + ); assert!(ParentKeys::::get(old_hotkey, netuid).is_empty()); // Verify ChildKeys swap - assert_eq!(ChildKeys::::get(new_hotkey, netuid), vec![(300u64, child1), (400u64, child2)]); + assert_eq!( + ChildKeys::::get(new_hotkey, netuid), + vec![(300u64, child1), (400u64, child2)] + ); assert!(ChildKeys::::get(old_hotkey, netuid).is_empty()); // Verify parent's ChildKeys update - assert_eq!(ChildKeys::::get(parent1, netuid), vec![(100u64, new_hotkey), (500u64, U256::from(8))]); - assert_eq!(ChildKeys::::get(parent2, netuid), vec![(200u64, new_hotkey), (600u64, U256::from(9))]); + assert_eq!( + ChildKeys::::get(parent1, netuid), + vec![(100u64, new_hotkey), (500u64, U256::from(8))] + ); + assert_eq!( + ChildKeys::::get(parent2, netuid), + vec![(200u64, new_hotkey), (600u64, U256::from(9))] + ); }); -} \ No newline at end of file +}