diff --git a/tests/e2e_tests/test_commit_weights.py b/tests/e2e_tests/test_commit_weights.py index 4d966c8c37..c4701ad71e 100644 --- a/tests/e2e_tests/test_commit_weights.py +++ b/tests/e2e_tests/test_commit_weights.py @@ -7,6 +7,7 @@ from tests.e2e_tests.utils.chain_interactions import ( sudo_set_admin_utils, sudo_set_hyperparameter_bool, + use_and_wait_for_next_nonce, wait_epoch, ) @@ -227,42 +228,44 @@ async def test_commit_weights_uses_next_nonce(local_chain, subtensor, alice_wall salt3[0] += 2 # Increment the first byte to produce a different commit hash # Commit all three salts - success, message = subtensor.commit_weights( - alice_wallet, - netuid, - salt=salt, - uids=weight_uids, - weights=weight_vals, - wait_for_inclusion=False, # Don't wait for inclusion, we are testing the nonce when there is a tx in the pool - wait_for_finalization=False, - ) - - assert success is True - - success, message = subtensor.commit_weights( - alice_wallet, - netuid, - salt=salt2, - uids=weight_uids, - weights=weight_vals, - wait_for_inclusion=False, - wait_for_finalization=False, - ) - - assert success is True - - # Commit the third salt - success, message = subtensor.commit_weights( - alice_wallet, - netuid, - salt=salt3, - uids=weight_uids, - weights=weight_vals, - wait_for_inclusion=False, - wait_for_finalization=False, - ) - - assert success is True + async with use_and_wait_for_next_nonce(subtensor, alice_wallet): + success, message = subtensor.commit_weights( + alice_wallet, + netuid, + salt=salt, + uids=weight_uids, + weights=weight_vals, + wait_for_inclusion=False, # Don't wait for inclusion, we are testing the nonce when there is a tx in the pool + wait_for_finalization=False, + ) + + assert success is True + + async with use_and_wait_for_next_nonce(subtensor, alice_wallet): + success, message = subtensor.commit_weights( + alice_wallet, + netuid, + salt=salt2, + uids=weight_uids, + weights=weight_vals, + wait_for_inclusion=False, + wait_for_finalization=False, + ) + + assert success is True + + async with use_and_wait_for_next_nonce(subtensor, alice_wallet): + success, message = subtensor.commit_weights( + alice_wallet, + netuid, + salt=salt3, + uids=weight_uids, + weights=weight_vals, + wait_for_inclusion=False, + wait_for_finalization=False, + ) + + assert success is True # Wait a few blocks await asyncio.sleep(10) # Wait for the txs to be included in the chain diff --git a/tests/e2e_tests/test_set_weights.py b/tests/e2e_tests/test_set_weights.py index 8a279e3ccf..8211e62aa9 100644 --- a/tests/e2e_tests/test_set_weights.py +++ b/tests/e2e_tests/test_set_weights.py @@ -6,6 +6,7 @@ from tests.e2e_tests.utils.chain_interactions import ( sudo_set_hyperparameter_bool, sudo_set_admin_utils, + use_and_wait_for_next_nonce, wait_epoch, ) @@ -108,16 +109,17 @@ async def test_set_weights_uses_next_nonce(local_chain, subtensor, alice_wallet) # Set weights for each subnet for netuid in netuids: - success, message = subtensor.set_weights( - alice_wallet, - netuid, - uids=weight_uids, - weights=weight_vals, - wait_for_inclusion=False, # Don't wait for inclusion, we are testing the nonce when there is a tx in the pool - wait_for_finalization=False, - ) - - assert success is True, f"Failed to set weights for subnet {netuid}" + async with use_and_wait_for_next_nonce(subtensor, alice_wallet): + success, message = subtensor.set_weights( + alice_wallet, + netuid, + uids=weight_uids, + weights=weight_vals, + wait_for_inclusion=False, # Don't wait for inclusion, we are testing the nonce when there is a tx in the pool + wait_for_finalization=False, + ) + + assert success is True, message # Wait for the txs to be included in the chain await wait_epoch(subtensor, netuid=netuids[-1], times=4) diff --git a/tests/e2e_tests/utils/chain_interactions.py b/tests/e2e_tests/utils/chain_interactions.py index 58529c4a3a..15998f93e0 100644 --- a/tests/e2e_tests/utils/chain_interactions.py +++ b/tests/e2e_tests/utils/chain_interactions.py @@ -4,6 +4,7 @@ """ import asyncio +import contextlib import unittest.mock from typing import Union, Optional, TYPE_CHECKING @@ -150,6 +151,30 @@ async def wait_interval( ) +@contextlib.asynccontextmanager +async def use_and_wait_for_next_nonce( + subtensor: "Subtensor", + wallet: "Wallet", + sleep: float = 0.25, + timeout: float = 15.0, +): + """ + ContextManager that makes sure the Nonce has been consumed after sending Extrinsic. + """ + + nonce = subtensor.substrate.get_account_next_index(wallet.hotkey.ss58_address) + + yield + + async def wait_for_new_nonce(): + while nonce == subtensor.substrate.get_account_next_index( + wallet.hotkey.ss58_address + ): + await asyncio.sleep(sleep) + + await asyncio.wait_for(wait_for_new_nonce(), timeout) + + # Helper to execute sudo wrapped calls on the chain def sudo_set_admin_utils( substrate: "SubstrateInterface",