Fix E2E Tests: wait for new nonce

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

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)

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",