fix test_dynamic_parent_child_relationships

Author: camfairchild

pallets/subtensor/src/tests/children.rs

@@ -3412,7 +3412,7 @@ fn test_parent_child_chain_emission() {
 // - Runs second epoch and distributes emissions
 // - Checks final emission distribution and stake updates
 // - Verifies correct parent-child relationships and stake proportions
-// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test children -- test_dynamic_parent_child_relationships --exact --nocapture
+// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --package pallet-subtensor --lib -- tests::children::test_dynamic_parent_child_relationships --exact --show-output
 #[test]
 fn test_dynamic_parent_child_relationships() {
     new_test_ext(1).execute_with(|| {
@@ -3432,14 +3432,66 @@ fn test_dynamic_parent_child_relationships() {
         register_ok_neuron(netuid, child1, coldkey_child1, 0);
         register_ok_neuron(netuid, child2, coldkey_child2, 0);
 
+        let chk_take_1 = SubtensorModule::get_childkey_take(&child1, netuid);
+        let chk_take_2 = SubtensorModule::get_childkey_take(&child2, netuid);
+        log::info!("child take 1: {:?}", chk_take_1);
+        log::info!("child take 2: {:?}", chk_take_2);
+
         // Add initial stakes
-        SubtensorModule::add_balance_to_coldkey_account(&coldkey_parent, 500_000);
-        SubtensorModule::add_balance_to_coldkey_account(&coldkey_child1, 50_000);
-        SubtensorModule::add_balance_to_coldkey_account(&coldkey_child2, 30_000);
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey_parent, 500_000 + 1_000);
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey_child1, 50_000 + 1_000);
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey_child2, 30_000 + 1_000);
+
+        // Swap to alpha
+        let total_tao: I96F32 = I96F32::from_num(500_000 + 50_000 + 30_000);
+        let total_alpha: I96F32 = I96F32::from_num(SubtensorModule::swap_tao_for_alpha(
+            netuid,
+            total_tao.saturating_to_num::<u64>(),
+        ));
+        log::info!("total_alpha: {:?}", total_alpha);
 
-        SubtensorModule::add_stake(RuntimeOrigin::signed(coldkey_parent), parent, netuid, 500_000).unwrap();
-        SubtensorModule::add_stake(RuntimeOrigin::signed(coldkey_child1), child1, netuid, 50_000).unwrap();
-        SubtensorModule::add_stake(RuntimeOrigin::signed(coldkey_child2), child2, netuid, 30_000).unwrap();
+        // Set the stakes directly
+        // This avoids needing to swap tao to alpha, impacting the initial stake distribution.
+        SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &parent,
+            &coldkey_parent,
+            netuid,
+            (total_alpha * I96F32::from_num(500_000) / total_tao).saturating_to_num::<u64>(),
+        );
+        SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &child1,
+            &coldkey_child1,
+            netuid,
+            (total_alpha * I96F32::from_num(50_000) / total_tao).saturating_to_num::<u64>(),
+        );
+        SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &child2,
+            &coldkey_child2,
+            netuid,
+            (total_alpha * I96F32::from_num(30_000) / total_tao).saturating_to_num::<u64>(),
+        );
+
+        // Get old stakes
+        let stake_parent_0: u64 = SubtensorModule::get_stake_for_hotkey_on_subnet(&parent, netuid);
+        let stake_child1_0: u64 = SubtensorModule::get_stake_for_hotkey_on_subnet(&child1, netuid);
+        let stake_child2_0: u64 = SubtensorModule::get_stake_for_hotkey_on_subnet(&child2, netuid);
+        log::info!("stake_parent_0: {:?}", stake_parent_0);
+        log::info!("stake_child1_0: {:?}", stake_child1_0);
+        log::info!("stake_child2_0: {:?}", stake_child2_0);
+
+        let total_stake_0: u64 = stake_parent_0 + stake_child1_0 + stake_child2_0;
+
+        // Assert initial stake is correct
+        let rel_stake_parent_0 = I96F32::from_num(stake_parent_0) / total_tao;
+        let rel_stake_child1_0 = I96F32::from_num(stake_child1_0) / total_tao;
+        let rel_stake_child2_0 = I96F32::from_num(stake_child2_0) / total_tao;
+
+        log::info!("rel_stake_parent_0: {:?}", rel_stake_parent_0);
+        log::info!("rel_stake_child1_0: {:?}", rel_stake_child1_0);
+        log::info!("rel_stake_child2_0: {:?}", rel_stake_child2_0);
+        assert_eq!(rel_stake_parent_0, I96F32::from_num(500_000) / total_tao);
+        assert_eq!(rel_stake_child1_0, I96F32::from_num(50_000) / total_tao);
+        assert_eq!(rel_stake_child2_0, I96F32::from_num(30_000) / total_tao);
 
         mock_set_children(&coldkey_parent, &parent, netuid, &[(u64::MAX / 2, child1)]);
 
@@ -3463,73 +3515,112 @@ fn test_dynamic_parent_child_relationships() {
         ));
 
         // Run first epoch
-        let hardcoded_emission: u64 = 1_000_000; // 1 million (adjust as needed)
+        let hardcoded_emission: u64 = 1_000_000;
 
         // Step blocks to allow for emission distribution
         step_block(11);
         step_rate_limit(&TransactionType::SetChildren, netuid);
 
+        // Get total stake after first payout
+        let total_stake_1 = SubtensorModule::get_stake_for_hotkey_on_subnet(&parent, netuid)
+            + SubtensorModule::get_stake_for_hotkey_on_subnet(&child1, netuid)
+            + SubtensorModule::get_stake_for_hotkey_on_subnet(&child2, netuid);
+        log::info!("total_stake_1: {:?}", total_stake_1);
+
         // Change parent-child relationships
-        mock_set_children(&coldkey_parent, &parent, netuid, &[(u64::MAX / 4, child1), (u64::MAX / 3, child2)]);
+        mock_set_children(
+            &coldkey_parent,
+            &parent,
+            netuid,
+            &[(u64::MAX / 4, child1), (u64::MAX / 3, child2)],
+        );
 
         // Step blocks again to allow for emission distribution
         step_block(11);
 
+        // Get total stake after second payout
+        let total_stake_2 = SubtensorModule::get_stake_for_hotkey_on_subnet(&parent, netuid)
+            + SubtensorModule::get_stake_for_hotkey_on_subnet(&child1, netuid)
+            + SubtensorModule::get_stake_for_hotkey_on_subnet(&child2, netuid);
+        log::info!("total_stake_2: {:?}", total_stake_2);
+
         // Check final emission distribution
-        let parent_stake: u64 = SubtensorModule::get_stake_for_hotkey_on_subnet(&parent, netuid);
-        let child1_stake: u64 = SubtensorModule::get_stake_for_hotkey_on_subnet(&child1, netuid);
-        let child2_stake: u64 = SubtensorModule::get_stake_for_hotkey_on_subnet(&child2, netuid);
+        let stake_parent_2: u64 =
+            SubtensorModule::get_inherited_for_hotkey_on_subnet(&parent, netuid);
+        let stake_child1_2: u64 =
+            SubtensorModule::get_inherited_for_hotkey_on_subnet(&child1, netuid);
+        let stake_child2_2: u64 =
+            SubtensorModule::get_inherited_for_hotkey_on_subnet(&child2, netuid);
+        let total_parent_stake = SubtensorModule::get_stake_for_hotkey_on_subnet(&parent, netuid);
+        let total_child1_stake = SubtensorModule::get_stake_for_hotkey_on_subnet(&child1, netuid);
+        let total_child2_stake = SubtensorModule::get_stake_for_hotkey_on_subnet(&child2, netuid);
 
         log::info!("Final stakes:");
-        log::info!("Parent stake: {}", parent_stake);
-        log::info!("Child1 stake: {}", child1_stake);
-        log::info!("Child2 stake: {}", child2_stake);
+        log::info!("Parent stake: {}", stake_parent_2);
+        log::info!("Child1 stake: {}", stake_child1_2);
+        log::info!("Child2 stake: {}", stake_child2_2);
+
+        // Payout 1
+        let payout_1 = total_stake_1 - total_stake_0;
+        log::info!("payout_1: {:?}", payout_1);
+
+        // Payout 2
+        let payout_2 = total_stake_2 - total_stake_1;
+        log::info!("payout_2: {:?}", payout_2);
+
+        let total_emission: I96F32 = I96F32::from_num(payout_1 + payout_2);
 
-        const TOLERANCE: u64 = 5; // Allow for a small discrepancy due to potential rounding
+        #[allow(non_snake_case)]
+        let TOLERANCE: I96F32 = I96F32::from_num(0.001); // Allow for a small discrepancy due to potential rounding
 
         // Precise assertions with tolerance
+        log::info!("total_emission: {:?}", total_emission);
+        let expected_parent_stake = ((I96F32::from_num(stake_parent_0)
+            + total_emission * rel_stake_parent_0)
+            * I96F32::from_num(5))
+        .saturating_div(I96F32::from_num(12));
         assert!(
-            (parent_stake as i128 - 926725i128).abs() <= TOLERANCE as i128,
-            "Parent stake should be close to 926,725, but was {}",
-            parent_stake
+            (I96F32::from_num(stake_parent_2) - expected_parent_stake).abs()
+                / expected_parent_stake
+                <= TOLERANCE,
+            "Parent stake should be close to {:?}, but was {}",
+            expected_parent_stake,
+            stake_parent_2
         );
         // Parent stake calculation:
         // Initial stake: 500,000
-        // First epoch: ~862,500 (500,000 + 725,000 * 1/2)
-        // Second epoch: ~926,725 (862,500 + 725,000 * 5/12)
+        // First epoch: 1/2 parent_stake
+        // Second epoch: 5/12 parent_stake
 
+        let expected_child1_stake = total_emission * rel_stake_child1_0
+            + I96F32::from_num(stake_child1_0 + (total_parent_stake) * 1 / 4);
         assert!(
-            (child1_stake as i64 - 778446).abs() <= TOLERANCE as i64,
-            "Child1 stake should be close to 778,446, but was {}",
-            child1_stake
+            (I96F32::from_num(stake_child1_2) - expected_child1_stake).abs()
+                / expected_child1_stake
+                <= TOLERANCE,
+            "Child1 stake should be close to {:?}, but was {}",
+            expected_child1_stake,
+            stake_child1_2
         );
         // Child1 stake calculation:
         // Initial stake: 50,000
-        // First epoch: ~412,500 (50,000 + 725,000 * 1/2)
-        // Second epoch: ~778,446 (412,500 + 725,000 * 1/2 * 1/4 + 137,500)
+        // First epoch: 1/2 parent_stake + child1_stake
+        // Second epoch: 1/4 parent_stake + child1_stake
 
+        let expected_child2_stake = total_emission * rel_stake_child2_0
+            + I96F32::from_num(stake_child2_0 + (total_parent_stake) * 1 / 3);
         assert!(
-            (child2_stake as i64 - 874826).abs() <= TOLERANCE as i64,
-            "Child2 stake should be close to 874,826, but was {}",
-            child2_stake
+            (I96F32::from_num(stake_child2_2) - expected_child2_stake).abs()
+                / expected_child2_stake
+                <= TOLERANCE,
+            "Child2 stake should be close to {:?}, but was {}",
+            expected_child2_stake,
+            stake_child2_2
         );
         // Child2 stake calculation:
         // Initial stake: 30,000
-        // First epoch: ~167,500 (30,000 + 137,500)
-        // Second epoch: ~874,826 (167,500 + 725,000 * 1/2 * 1/3 + 137,500)
-
-        // Check that the total stake has increased by approximately twice the hardcoded emission amount
-        let total_stake: u64 = parent_stake + child1_stake + child2_stake;
-        let initial_total_stake: u64 = 500_000 + 50_000 + 30_000;
-        let total_emission: u64 = 2 * hardcoded_emission;
-        assert!(
-            (total_stake as i64 - (initial_total_stake + total_emission) as i64).abs() <= TOLERANCE as i64,
-            "Total stake should have increased by approximately twice the hardcoded emission amount"
-        );
-        // Total stake calculation:
-        // Initial total stake: 500,000 + 50,000 + 30,000 = 580,000
-        // Total emission: 2 * 1,000,000 = 2,000,000
-        // Expected total stake: 580,000 + 2,000,000 = 2,580,000
+        // First epoch: child2_stake
+        // Second epoch: 1/3 parent_stake + child2_stake
 
         // Additional checks for parent-child relationships
         let parent_children: Vec<(u64, U256)> = SubtensorModule::get_children(&parent, netuid);
@@ -3562,21 +3653,9 @@ fn test_dynamic_parent_child_relationships() {
 
         // Check that child2 has received more stake than child1
         assert!(
-            child2_stake > child1_stake,
+            stake_child2_2 > stake_child1_2,
             "Child2 should have received more emission than Child1 due to higher proportion"
         );
         // Child2 stake (874,826) > Child1 stake (778,446)
-
-        // Check the approximate difference between child2 and child1 stakes
-        let stake_difference: u64 = child2_stake - child1_stake;
-        assert!(
-            (stake_difference as i64 - 96_380).abs() <= TOLERANCE as i64,
-            "The difference between Child2 and Child1 stakes should be close to 96,380, but was {}",
-            stake_difference
-        );
-        // Stake difference calculation:
-        // Child2 stake: 874,826
-        // Child1 stake: 778,446
-        // Difference: 874,826 - 778,446 = 96,380
     });
 }