Merge pull request #3153 from opentensor/fix/roman/root-claim-e2e-tests

[v10] Bringing RootClam tests to life

Author: basfroman

tests/e2e_tests/test_root_claim.py

@@ -1,5 +1,9 @@
 from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
+from tests.e2e_tests.utils.set_subnet_moving_price import (
+    increase_subnet_ema,
+    async_increase_subnet_ema,
+)
 from tests.e2e_tests.utils import (
     AdminUtils,
     NETUID,
@@ -17,7 +21,6 @@
 PROOF_COUNTER = 2
 
 
-@pytest.mark.skip(reason="we need to figure out how to emulate behavior")
 def test_root_claim_swap(subtensor, alice_wallet, bob_wallet, charlie_wallet):
     """Tests root claim Swap logic.
 
@@ -51,6 +54,9 @@ def test_root_claim_swap(subtensor, alice_wallet, bob_wallet, charlie_wallet):
         ]
     )
 
+    # Here is the damn magic with EMA
+    assert increase_subnet_ema(subtensor=subtensor, sudo_wallet=alice_wallet)
+
     stake_balance = Balance.from_tao(10)
 
     # Stake to Alice in ROOT
@@ -112,7 +118,6 @@ def test_root_claim_swap(subtensor, alice_wallet, bob_wallet, charlie_wallet):
         proof_counter -= 1
 
 
-@pytest.mark.skip(reason="we need to figure out how to emulate behavior")
 @pytest.mark.asyncio
 async def test_root_claim_swap_async(
     async_subtensor, alice_wallet, bob_wallet, charlie_wallet
@@ -149,6 +154,11 @@ async def test_root_claim_swap_async(
         ]
     )
 
+    # Here is the damn magic with EMA
+    assert await async_increase_subnet_ema(
+        subtensor=async_subtensor, sudo_wallet=alice_wallet
+    )
+
     stake_balance = Balance.from_tao(10)
 
     # Stake to Alice in ROOT
@@ -214,7 +224,6 @@ async def test_root_claim_swap_async(
         proof_counter -= 1
 
 
-@pytest.mark.skip(reason="we need to figure out how to emulate behavior")
 def test_root_claim_keep_with_zero_num_root_auto_claims(
     subtensor, alice_wallet, bob_wallet, charlie_wallet
 ):
@@ -280,6 +289,9 @@ def test_root_claim_keep_with_zero_num_root_auto_claims(
         == "Keep"
     )
 
+    # Here is the damn magic with EMA
+    assert increase_subnet_ema(subtensor=subtensor, sudo_wallet=alice_wallet)
+
     stake_balance = Balance.from_tao(1000)  # just a dream - stake 1000 TAO to SN0 :D
 
     # Stake from Charlie to Alice in ROOT
@@ -376,7 +388,6 @@ def test_root_claim_keep_with_zero_num_root_auto_claims(
     logging.console.info(f"SN2 Stake: {stake_after_charlie}")
 
 
-@pytest.mark.skip(reason="we need to figure out how to emulate behavior")
 @pytest.mark.asyncio
 async def test_root_claim_keep_with_zero_num_root_auto_claims_async(
     async_subtensor, alice_wallet, bob_wallet, charlie_wallet
@@ -445,6 +456,11 @@ async def test_root_claim_keep_with_zero_num_root_auto_claims_async(
         == "Keep"
     )
 
+    # Here is the damn magic with EMA
+    assert await async_increase_subnet_ema(
+        subtensor=async_subtensor, sudo_wallet=alice_wallet
+    )
+
     stake_balance = Balance.from_tao(1000)  # just a dream - stake 1000 TAO to SN0 :D
 
     # Stake from Charlie to Alice in ROOT
@@ -547,7 +563,6 @@ async def test_root_claim_keep_with_zero_num_root_auto_claims_async(
     logging.console.info(f"SN2 Stake: {stake_after_charlie}")
 
 
-@pytest.mark.skip(reason="we need to figure out how to emulate behavior")
 def test_root_claim_keep_with_random_auto_claims(
     subtensor, alice_wallet, bob_wallet, charlie_wallet, dave_wallet
 ):
@@ -610,7 +625,10 @@ def test_root_claim_keep_with_random_auto_claims(
         == "Keep"
     )
 
-    stake_balance = Balance.from_tao(1000)  # just a dream - stake 1000 TAO to SN0 :D
+    # Set EMA to enable root_sell_flag
+    assert increase_subnet_ema(subtensor=subtensor, sudo_wallet=alice_wallet)
+
+    stake_balance = Balance.from_tao(1000)
 
     # Stake from Charlie to Alice in ROOT
     response = subtensor.staking.add_stake(
@@ -621,6 +639,13 @@ def test_root_claim_keep_with_random_auto_claims(
     )
     assert response.success, response.message
 
+    # Skip the epoch in which stake was installed. Emission doesn't occur in the same epoch as stake installation
+    logging.console.info("Skipping stake epoch")
+    next_epoch_start_block = subtensor.subnets.get_next_epoch_start_block(
+        netuid=root_sn.netuid
+    )
+    subtensor.wait_for_block(next_epoch_start_block)
+
     proof_counter = PROOF_COUNTER
 
     prev_claimed_stake_charlie = subtensor.staking.get_stake(
@@ -635,8 +660,7 @@ def test_root_claim_keep_with_random_auto_claims(
         netuid=sn2.netuid,
     )
 
-    # proof that ROOT stake increases each epoch even RootClaimType is Keep bc of random auto claim takes min 5 coldkeys
-    # to do release emissions
+    # Wait for epochs and check that stake increases due to random auto claims
     while proof_counter > 0:
         next_epoch_start_block = subtensor.subnets.get_next_epoch_start_block(
             root_sn.netuid
@@ -649,21 +673,24 @@ def test_root_claim_keep_with_random_auto_claims(
             hotkey_ss58=alice_wallet.hotkey.ss58_address,
             netuid=sn2.netuid,
         )
-        assert claimed_stake_charlie > prev_claimed_stake_charlie
+        assert claimed_stake_charlie > prev_claimed_stake_charlie, (
+            f"Stake did not increase: {claimed_stake_charlie} <= {prev_claimed_stake_charlie}"
+        )
         prev_claimed_stake_charlie = claimed_stake_charlie
 
         root_claimed_charlie = subtensor.staking.get_root_claimed(
             coldkey_ss58=charlie_wallet.coldkey.ss58_address,
             hotkey_ss58=alice_wallet.hotkey.ss58_address,
             netuid=sn2.netuid,
         )
-        assert root_claimed_charlie > prev_root_claimed_charlie
+        assert root_claimed_charlie > prev_root_claimed_charlie, (
+            f"Root claimed did not increase: {root_claimed_charlie} <= {prev_root_claimed_charlie}"
+        )
         prev_root_claimed_charlie = root_claimed_charlie
 
         proof_counter -= 1
 
 
-@pytest.mark.skip(reason="we need to figure out how to emulate behavior")
 @pytest.mark.asyncio
 async def test_root_claim_keep_with_random_auto_claims_async(
     async_subtensor, alice_wallet, bob_wallet, charlie_wallet, dave_wallet
@@ -729,6 +756,11 @@ async def test_root_claim_keep_with_random_auto_claims_async(
         == "Keep"
     )
 
+    # Here is the damn magic with EMA
+    assert await async_increase_subnet_ema(
+        subtensor=async_subtensor, sudo_wallet=alice_wallet
+    )
+
     stake_balance = Balance.from_tao(1000)  # just a dream - stake 1000 TAO to SN0 :D
 
     # Stake from Charlie to Alice in ROOT
@@ -740,6 +772,13 @@ async def test_root_claim_keep_with_random_auto_claims_async(
     )
     assert response.success, response.message
 
+    # Skip the epoch in which stake was installed. Emission doesn't occur in the same epoch as stake installation
+    logging.console.info("Skipping stake epoch")
+    next_epoch_start_block = await async_subtensor.subnets.get_next_epoch_start_block(
+        netuid=root_sn.netuid
+    )
+    await async_subtensor.wait_for_block(next_epoch_start_block)
+
     proof_counter = PROOF_COUNTER
 
     prev_claimed_stake_charlie = await async_subtensor.staking.get_stake(
@@ -754,8 +793,7 @@ async def test_root_claim_keep_with_random_auto_claims_async(
         netuid=sn2.netuid,
     )
 
-    # proof that ROOT stake increases each epoch even RootClaimType is Keep bc of random auto claim takes min 5 coldkeys
-    # to do release emissions
+    # Wait for epochs and check that stake increases due to random auto claims
     while proof_counter > 0:
         next_epoch_start_block = (
             await async_subtensor.subnets.get_next_epoch_start_block(root_sn.netuid)
@@ -768,15 +806,19 @@ async def test_root_claim_keep_with_random_auto_claims_async(
             hotkey_ss58=alice_wallet.hotkey.ss58_address,
             netuid=sn2.netuid,
         )
-        assert claimed_stake_charlie > prev_claimed_stake_charlie
+        assert claimed_stake_charlie > prev_claimed_stake_charlie, (
+            f"Stake did not increase: {claimed_stake_charlie} <= {prev_claimed_stake_charlie}"
+        )
         prev_claimed_stake_charlie = claimed_stake_charlie
 
         root_claimed_charlie = await async_subtensor.staking.get_root_claimed(
             coldkey_ss58=charlie_wallet.coldkey.ss58_address,
             hotkey_ss58=alice_wallet.hotkey.ss58_address,
             netuid=sn2.netuid,
         )
-        assert root_claimed_charlie > prev_root_claimed_charlie
+        assert root_claimed_charlie > prev_root_claimed_charlie, (
+            f"Root claimed did not increase: {root_claimed_charlie} <= {prev_root_claimed_charlie}"
+        )
         prev_root_claimed_charlie = root_claimed_charlie
 
         proof_counter -= 1

tests/e2e_tests/utils/set_subnet_moving_price.py

@@ -0,0 +1,121 @@
+from bittensor_wallet import Wallet
+
+from bittensor.core.extrinsics.utils import sudo_call_extrinsic
+from bittensor.core.extrinsics.asyncex.utils import (
+    sudo_call_extrinsic as async_sudo_call_extrinsic,
+)
+from bittensor.extras.subtensor_api import SubtensorApi
+
+
+def check_root_sell_flag(subtensor: "SubtensorApi") -> tuple[bool, int]:
+    """Checks the current state of root_sell_flag."""
+    # Get all non-root subnets
+    subnets = [
+        netuid for netuid in subtensor.subnets.get_all_subnets_netuid() if netuid != 0
+    ]
+    total_ema = 0
+    for netuid in subnets:
+        # Get moving_alpha_price through subnet's Dynamic info data
+        moving_price = subtensor.subnets.subnet(netuid).moving_price
+        total_ema += moving_price
+    root_sell_flag = total_ema > 1.0
+    return root_sell_flag, total_ema
+
+
+async def async_check_root_sell_flag(subtensor: "SubtensorApi") -> tuple[bool, int]:
+    """Checks the current state of root_sell_flag."""
+    # Get all non-root subnets
+    subnets = [
+        netuid
+        for netuid in await subtensor.subnets.get_all_subnets_netuid()
+        if netuid != 0
+    ]
+    total_ema = 0
+    for netuid in subnets:
+        # Get moving_alpha_price through subnet's Dynamic info data
+        moving_price = (await subtensor.subnets.subnet(netuid)).moving_price
+        total_ema += moving_price
+    root_sell_flag = total_ema > 1.0
+    return root_sell_flag, total_ema
+
+
+def set_huge_ema_before_test(subtensor: "SubtensorApi", sudo_wallet: "Wallet"):
+    """Sets huge EMA value into SubnetMovingPrice storage."""
+
+    encoded_storage_key = (
+        "0x658faa385070e074c85bf6b568cf05551abf1b0f4fd14f7b72ee50f9d91d59150200"
+    )
+    encoded_storage_value = b"003665c4ffc99a3b0000000000000000"
+    items = [encoded_storage_key, encoded_storage_value]
+
+    return sudo_call_extrinsic(
+        subtensor=subtensor.inner_subtensor,
+        wallet=sudo_wallet,
+        call_function="set_storage",
+        call_params={"items": [items]},
+        call_module="System",
+    )
+
+
+async def async_set_huge_ema_before_test(
+    subtensor: "SubtensorApi", sudo_wallet: "Wallet"
+):
+    """Sets huge EMA value into SubnetMovingPrice storage."""
+
+    encoded_storage_key = (
+        "0x658faa385070e074c85bf6b568cf05551abf1b0f4fd14f7b72ee50f9d91d59150200"
+    )
+    encoded_storage_value = b"003665c4ffc99a3b0000000000000000"
+    items = [encoded_storage_key, encoded_storage_value]
+
+    return await async_sudo_call_extrinsic(
+        subtensor=subtensor.inner_subtensor,
+        wallet=sudo_wallet,
+        call_function="set_storage",
+        call_params={"items": [items]},
+        call_module="System",
+    )
+
+
+def increase_subnet_ema(subtensor: "SubtensorApi", sudo_wallet: "Wallet") -> bool:
+    """Increases EMA value in SubnetMovingPrice storage."""
+    try:
+        root_sell_flag, sns_sum_of_ema = check_root_sell_flag(subtensor)
+        assert not root_sell_flag
+        assert sns_sum_of_ema < 1.0
+
+        response = set_huge_ema_before_test(
+            subtensor=subtensor, sudo_wallet=sudo_wallet
+        )
+        assert response.success, response.message
+        subtensor.wait_for_block()
+
+        root_sell_flag, sns_sum_of_ema = check_root_sell_flag(subtensor)
+        assert root_sell_flag, "Root sell still false"
+        assert sns_sum_of_ema > 1.0, "SNs EMA sum wasn't increased greater than 1.0"
+        return True
+    except AssertionError:
+        return False
+
+
+async def async_increase_subnet_ema(
+    subtensor: "SubtensorApi", sudo_wallet: "Wallet"
+) -> bool:
+    """Increases EMA value in SubnetMovingPrice storage."""
+    try:
+        root_sell_flag, sns_sum_of_ema = await async_check_root_sell_flag(subtensor)
+        assert not root_sell_flag
+        assert sns_sum_of_ema < 1.0
+
+        response = await async_set_huge_ema_before_test(
+            subtensor=subtensor, sudo_wallet=sudo_wallet
+        )
+        assert response.success, response.message
+        await subtensor.wait_for_block()
+
+        root_sell_flag, sns_sum_of_ema = await async_check_root_sell_flag(subtensor)
+        assert root_sell_flag, "Root sell still false"
+        assert sns_sum_of_ema > 1.0, "SNs EMA sum wasn't increased greater than 1.0"
+        return True
+    except AssertionError:
+        return False