Fix/e2e tests for stake lock owner alpha

bittensor/core/extrinsics/children.py

@@ -111,7 +111,7 @@ def root_set_pending_childkey_cooldown_extrinsic(
     period: Optional[int] = None,
     raise_error: bool = False,
     wait_for_inclusion: bool = True,
-    wait_for_finalization: bool = False,
+    wait_for_finalization: bool = True,
     wait_for_revealed_execution: bool = True,
 ) -> ExtrinsicResponse:
     """

tests/e2e_tests/test_dendrite.py

@@ -46,7 +46,7 @@ def test_dendrite(subtensor, templates, alice_wallet, bob_wallet):
         REGISTER_SUBNET(alice_wallet),
         SUDO_SET_TEMPO(alice_wallet, AdminUtils, True, NETUID, TEMPO_TO_SET),
         ACTIVATE_SUBNET(alice_wallet),
-        SUDO_SET_MAX_ALLOWED_VALIDATORS(alice_wallet, AdminUtils, True, NETUID, 1),
+        SUDO_SET_MAX_ALLOWED_VALIDATORS(alice_wallet, AdminUtils, True, NETUID, 2),
         SUDO_SET_WEIGHTS_SET_RATE_LIMIT(alice_wallet, AdminUtils, True, NETUID, 10),
         REGISTER_NEURON(bob_wallet),
     ]
@@ -160,7 +160,7 @@ async def test_dendrite_async(async_subtensor, templates, alice_wallet, bob_wall
         REGISTER_SUBNET(alice_wallet),
         SUDO_SET_TEMPO(alice_wallet, AdminUtils, True, NETUID, TEMPO_TO_SET),
         ACTIVATE_SUBNET(alice_wallet),
-        SUDO_SET_MAX_ALLOWED_VALIDATORS(alice_wallet, AdminUtils, True, NETUID, 1),
+        SUDO_SET_MAX_ALLOWED_VALIDATORS(alice_wallet, AdminUtils, True, NETUID, 2),
         SUDO_SET_WEIGHTS_SET_RATE_LIMIT(alice_wallet, AdminUtils, True, NETUID, 10),
         REGISTER_NEURON(bob_wallet),
     ]

tests/e2e_tests/test_metagraph.py

@@ -1,6 +1,7 @@
 import os.path
 import shutil
 import time
+from dataclasses import replace
 import numpy as np
 import pytest
 
@@ -27,6 +28,12 @@
 NULL_KEY = tuple(bytearray(32))
 
 
+def _strip_unique_fields_for_metagraph_parity(m: MetagraphInfo) -> MetagraphInfo:
+    # get_metagraph_info (selective mechagraph) contains validators/commitments;
+    # get_all_metagraphs_info does not get these fields from the chain.
+    return replace(m, validators=None, commitments=None)
+
+
 torch = LazyLoadedTorch()
 
 
@@ -833,7 +840,9 @@ def test_metagraph_info(subtensor, alice_wallet, bob_wallet):
     metagraph_infos = subtensor.metagraphs.get_all_metagraphs_info(block=block)
 
     assert len(metagraph_infos) == 4
-    assert metagraph_infos[-1] == metagraph_info
+    assert metagraph_infos[-1] == _strip_unique_fields_for_metagraph_parity(
+        metagraph_info
+    )
 
     # non-existed subnet
     metagraph_info = subtensor.metagraphs.get_metagraph_info(netuid=bob_sn.netuid + 1)
@@ -1088,7 +1097,9 @@ async def test_metagraph_info_async(async_subtensor, alice_wallet, bob_wallet):
     )
 
     assert len(metagraph_infos) == 4
-    assert metagraph_infos[-1] == metagraph_info
+    assert metagraph_infos[-1] == _strip_unique_fields_for_metagraph_parity(
+        metagraph_info
+    )
 
     # non-existed subnet
     metagraph_info = await async_subtensor.metagraphs.get_metagraph_info(

tests/e2e_tests/test_staking.py

@@ -20,6 +20,32 @@
 from tests.helpers.helpers import CloseInValue
 
 
+def _get_expected_balance_after_neuron_registrations(
+    initial_balance: Balance, sns: list[TestSubnet]
+) -> Balance:
+    """
+    Get the expected balance after neuron registrations
+    by subtracting the total registration cost from the initial balance.
+    """
+    total_registration_cost = Balance.from_rao(0)
+
+    for sn in sns:
+        for call in sn.calls:
+            if call.operation != REGISTER_NEURON.__name__:
+                continue
+
+            balance_before = call.response.data.get("balance_before")
+            balance_after = call.response.data.get("balance_after")
+
+            assert balance_before is not None and balance_after is not None, (
+                "REGISTER_NEURON response missing balance data."
+            )
+
+            total_registration_cost += balance_before - balance_after
+
+    return initial_balance - total_registration_cost
+
+
 def test_single_operation(subtensor, alice_wallet, bob_wallet):
     """
     Tests:
@@ -334,6 +360,9 @@ def test_batch_operations(subtensor, alice_wallet, bob_wallet):
     - Checks Accounts Balance
     """
     subnets_tested = 2
+    bob_balance_before_setup = subtensor.wallets.get_balance(
+        bob_wallet.coldkey.ss58_address
+    )
 
     sns = [TestSubnet(subtensor) for _ in range(subnets_tested)]
 
@@ -357,6 +386,9 @@ def test_batch_operations(subtensor, alice_wallet, bob_wallet):
         )
 
     netuids = [sn.netuid for sn in sns]
+    expected_bob_balance = _get_expected_balance_after_neuron_registrations(
+        bob_balance_before_setup, sns
+    )
 
     balances = subtensor.wallets.get_balances(
         alice_wallet.coldkey.ss58_address,
@@ -443,10 +475,7 @@ def test_batch_operations(subtensor, alice_wallet, bob_wallet):
         bob_wallet.coldkey.ss58_address,
     )
 
-    assert CloseInValue(  # Make sure we are within 0.0002 TAO due to tx fees
-        balances[bob_wallet.coldkey.ss58_address], Balance.from_rao(5_000_000)
-    ) == Balance.from_tao(999_999.7979)
-
+    assert balances[bob_wallet.coldkey.ss58_address] == expected_bob_balance
     assert balances[alice_wallet.coldkey.ss58_address] > alice_balance
 
 
@@ -460,6 +489,9 @@ async def test_batch_operations_async(async_subtensor, alice_wallet, bob_wallet)
     - Checks Accounts Balance
     """
     subnets_tested = 2
+    bob_balance_before_setup = await async_subtensor.wallets.get_balance(
+        bob_wallet.coldkey.ss58_address
+    )
 
     sns = [TestSubnet(async_subtensor) for _ in range(subnets_tested)]
 
@@ -483,6 +515,9 @@ async def test_batch_operations_async(async_subtensor, alice_wallet, bob_wallet)
         )
 
     netuids = [sn.netuid for sn in sns]
+    expected_bob_balance = _get_expected_balance_after_neuron_registrations(
+        bob_balance_before_setup, sns
+    )
 
     balances = await async_subtensor.wallets.get_balances(
         alice_wallet.coldkey.ss58_address,
@@ -569,10 +604,7 @@ async def test_batch_operations_async(async_subtensor, alice_wallet, bob_wallet)
         bob_wallet.coldkey.ss58_address,
     )
 
-    assert CloseInValue(  # Make sure we are within 0.0002 TAO due to tx fees
-        balances[bob_wallet.coldkey.ss58_address], Balance.from_rao(5_000_000)
-    ) == Balance.from_tao(999_999.7979)
-
+    assert balances[bob_wallet.coldkey.ss58_address] == expected_bob_balance
     assert balances[alice_wallet.coldkey.ss58_address] > alice_balance
 
 

tests/e2e_tests/test_subtensor_functions.py

@@ -15,6 +15,7 @@
     ACTIVATE_SUBNET,
     REGISTER_NEURON,
 )
+from tests.helpers import CloseInValue
 
 """
 Verifies:
@@ -142,12 +143,9 @@ def test_subtensor_extrinsics(subtensor, templates, alice_wallet, bob_wallet):
 
     bob_balance = subtensor.wallets.get_balance(bob_wallet.coldkeypub.ss58_address)
 
-    alice_sn.execute_steps(
-        [
-            ACTIVATE_SUBNET(alice_wallet),
-            REGISTER_NEURON(bob_wallet),
-        ]
-    )
+    alice_sn.execute_one(ACTIVATE_SUBNET(alice_wallet))
+    recycle_amount = subtensor.subnets.recycle(netuid)
+    reg_response = alice_sn.execute_one(REGISTER_NEURON(bob_wallet))
 
     # Verify Bob's UID on netuid 2 is 1
     assert (
@@ -157,17 +155,16 @@ def test_subtensor_extrinsics(subtensor, templates, alice_wallet, bob_wallet):
         == 1
     ), "UID for Bob's hotkey on netuid 2 is not 1 as expected"
 
-    # Fetch recycle_amount to register to the subnet
-    recycle_amount = subtensor.subnets.recycle(netuid)
-    fee = alice_sn.calls[-1].response.extrinsic_fee
+    fee = reg_response.extrinsic_fee
     bob_balance_post_reg = subtensor.wallets.get_balance(
         bob_wallet.coldkeypub.ss58_address
     )
 
-    # Ensure recycled amount is only deducted from the balance after registration
-    assert bob_balance - recycle_amount - fee == bob_balance_post_reg, (
-        "Balance for Bob is not correct after burned register"
-    )
+    # Burn decays every block and bumps after registration; hence tolerance
+    expected_post = bob_balance - recycle_amount - fee
+    assert (
+        CloseInValue(bob_balance_post_reg, Balance.from_tao(0.002)) == expected_post
+    ), "Balance for Bob is not correct after burned register"
 
     with templates.validator(alice_wallet, netuid):
         # wait for 5 seconds for the metagraph and subtensor to refresh with latest data
@@ -311,12 +308,9 @@ async def test_subtensor_extrinsics_async(
         bob_wallet.coldkeypub.ss58_address
     )
 
-    await alice_sn.async_execute_steps(
-        [
-            ACTIVATE_SUBNET(alice_wallet),
-            REGISTER_NEURON(bob_wallet),
-        ]
-    )
+    await alice_sn.async_execute_one(ACTIVATE_SUBNET(alice_wallet))
+    recycle_amount = await async_subtensor.subnets.recycle(netuid)
+    reg_response = await alice_sn.async_execute_one(REGISTER_NEURON(bob_wallet))
 
     # Verify Bob's UID on netuid 2 is 1
     assert (
@@ -326,17 +320,16 @@ async def test_subtensor_extrinsics_async(
         == 1
     ), "UID for Bob's hotkey on netuid 2 is not 1 as expected."
 
-    # Fetch recycle_amount to register to the subnet
-    recycle_amount = await async_subtensor.subnets.recycle(netuid)
-    fee = alice_sn.calls[-1].response.extrinsic_fee
+    fee = reg_response.extrinsic_fee
     bob_balance_post_reg = await async_subtensor.wallets.get_balance(
         bob_wallet.coldkeypub.ss58_address
     )
 
-    # Ensure recycled amount is only deducted from the balance after registration
-    assert bob_balance - recycle_amount - fee == bob_balance_post_reg, (
-        "Balance for Bob is not correct after burned register"
-    )
+    expected_post = bob_balance - recycle_amount - fee
+    # Burn decays every block and bumps after registration; hence tolerance
+    assert (
+        CloseInValue(bob_balance_post_reg, Balance.from_tao(0.005)) == expected_post
+    ), "Balance for Bob is not correct after burned register"
 
     # neuron_info_old = subtensor.get_neuron_for_pubkey_and_subnet(
     #     alice_wallet.hotkey.ss58_address, netuid=netuid

tests/e2e_tests/utils/e2e_test_utils.py

@@ -64,7 +64,7 @@ def clone_or_update_templates(specific_commit=None):
     """
     install_dir = template_path
     repo_mapping = {
-        templates_repo: "https://github.com/opentensor/subnet-template.git",
+        templates_repo: "https://github.com/latent-to/subnet-template.git",
     }
 
     cwd = os.getcwd()

tests/unit_tests/extrinsics/test_children.py

@@ -91,7 +91,7 @@ def test_root_set_pending_childkey_cooldown_extrinsic(subtensor, mocker, fake_wa
         period=None,
         raise_error=False,
         wait_for_inclusion=True,
-        wait_for_finalization=False,
+        wait_for_finalization=True,
     )
     assert success is True
     assert "Success" in message