tests: wait for Miner/Validator to fully start

Author: zyzniewski-reef

tests/e2e_tests/test_incentive.py

@@ -41,6 +41,9 @@ async def test_incentive(local_chain, subtensor, templates, alice_wallet, bob_wa
         len(subtensor.neurons(netuid=netuid)) == 2
     ), "Alice & Bob not registered in the subnet"
 
+    # Wait for the first epoch to pass
+    await wait_epoch(subtensor, netuid)
+
     # Get latest metagraph
     metagraph = subtensor.metagraph(netuid)
 
@@ -51,6 +54,9 @@ async def test_incentive(local_chain, subtensor, templates, alice_wallet, bob_wa
     assert alice_neuron.dividends == 0
     assert alice_neuron.stake.tao > 0
     assert alice_neuron.validator_trust == 0
+    assert alice_neuron.incentive == 0
+    assert alice_neuron.consensus == 0
+    assert alice_neuron.rank == 0
 
     bob_neuron = metagraph.neurons[1]
 
@@ -69,9 +75,10 @@ async def test_incentive(local_chain, subtensor, templates, alice_wallet, bob_wa
     )
 
     async with templates.miner(bob_wallet, netuid):
-        async with templates.validator(alice_wallet, netuid):
+        async with templates.validator(alice_wallet, netuid) as validator:
             # wait for the Validator to process and set_weights
-            await asyncio.sleep(5)
+            async with asyncio.timeout(15):
+                await validator.set_weights.wait()
 
             # Wait few epochs
             await wait_epoch(subtensor, netuid, times=4)
@@ -85,12 +92,15 @@ async def test_incentive(local_chain, subtensor, templates, alice_wallet, bob_wa
     assert alice_neuron.validator_permit is True
     assert alice_neuron.dividends == 1.0
     assert alice_neuron.stake.tao > 0
-    assert alice_neuron.validator_trust == 1
+    assert alice_neuron.validator_trust > 0.99
+    assert alice_neuron.incentive < 0.5
+    assert alice_neuron.consensus < 0.5
+    assert alice_neuron.rank < 0.5
 
     bob_neuron = metagraph.neurons[1]
-    assert bob_neuron.incentive == 1
-    assert bob_neuron.consensus == 1
-    assert bob_neuron.rank == 1
+    assert bob_neuron.incentive > 0.5
+    assert bob_neuron.consensus > 0.5
+    assert bob_neuron.rank > 0.5
     assert bob_neuron.trust == 1
 
     print("✅ Passed test_incentive")

tests/e2e_tests/utils/e2e_test_utils.py

@@ -1,5 +1,4 @@
 import asyncio
-import contextlib
 import os
 import shutil
 import subprocess
@@ -84,6 +83,107 @@ def uninstall_templates(install_dir):
 
 
 class Templates:
+    class Miner:
+        def __init__(self, dir, wallet, netuid):
+            self.dir = dir
+            self.wallet = wallet
+            self.netuid = netuid
+            self.process = None
+
+            self.started = asyncio.Event()
+
+        async def __aenter__(self):
+            self.process = await asyncio.create_subprocess_exec(
+                sys.executable,
+                f"{self.dir}/miner.py",
+                "--netuid",
+                str(self.netuid),
+                "--subtensor.network",
+                "local",
+                "--subtensor.chain_endpoint",
+                "ws://localhost:9944",
+                "--wallet.path",
+                self.wallet.path,
+                "--wallet.name",
+                self.wallet.name,
+                "--wallet.hotkey",
+                "default",
+                env={
+                    "BT_LOGGING_INFO": "1",
+                },
+                stdout=asyncio.subprocess.PIPE,
+            )
+
+            self.__reader_task = asyncio.create_task(self._reader())
+
+            async with asyncio.timeout(30):
+                await self.started.wait()
+
+            return self
+
+        async def __aexit__(self, exc_type, exc_value, traceback):
+            self.process.terminate()
+            self.__reader_task.cancel()
+
+            await self.process.wait()
+
+        async def _reader(self):
+            async for line in self.process.stdout:
+                if b"Starting main loop" in line:
+                    self.started.set()
+
+    class Validator:
+        def __init__(self, dir, wallet, netuid):
+            self.dir = dir
+            self.wallet = wallet
+            self.netuid = netuid
+            self.process = None
+
+            self.started = asyncio.Event()
+            self.set_weights = asyncio.Event()
+
+        async def __aenter__(self):
+            self.process = await asyncio.create_subprocess_exec(
+                sys.executable,
+                f"{self.dir}/validator.py",
+                "--netuid",
+                str(self.netuid),
+                "--subtensor.network",
+                "local",
+                "--subtensor.chain_endpoint",
+                "ws://localhost:9944",
+                "--wallet.path",
+                self.wallet.path,
+                "--wallet.name",
+                self.wallet.name,
+                "--wallet.hotkey",
+                "default",
+                env={
+                    "BT_LOGGING_INFO": "1",
+                },
+                stdout=asyncio.subprocess.PIPE,
+            )
+
+            self.__reader_task = asyncio.create_task(self._reader())
+
+            async with asyncio.timeout(30):
+                await self.started.wait()
+
+            return self
+
+        async def __aexit__(self, exc_type, exc_value, traceback):
+            self.process.terminate()
+            self.__reader_task.cancel()
+
+            await self.process.wait()
+
+        async def _reader(self):
+            async for line in self.process.stdout:
+                if b"Starting validator loop." in line:
+                    self.started.set()
+                elif b"Successfully set weights and Finalized." in line:
+                    self.set_weights.set()
+
     def __init__(self):
         self.dir = clone_or_update_templates()
 
@@ -93,52 +193,8 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, traceback):
         uninstall_templates(self.dir)
 
-    @contextlib.asynccontextmanager
-    async def miner(self, wallet, netuid):
-        process = await asyncio.create_subprocess_exec(
-            sys.executable,
-            f"{self.dir}/miner.py",
-            "--netuid",
-            str(netuid),
-            "--subtensor.network",
-            "local",
-            "--subtensor.chain_endpoint",
-            "ws://localhost:9944",
-            "--wallet.path",
-            wallet.path,
-            "--wallet.name",
-            wallet.name,
-            "--wallet.hotkey",
-            "default",
-        )
-
-        yield
-
-        process.terminate()
-
-        await process.wait()
-
-    @contextlib.asynccontextmanager
-    async def validator(self, wallet, netuid):
-        process = await asyncio.create_subprocess_exec(
-            sys.executable,
-            f"{self.dir}/validator.py",
-            "--netuid",
-            str(netuid),
-            "--subtensor.network",
-            "local",
-            "--subtensor.chain_endpoint",
-            "ws://localhost:9944",
-            "--wallet.path",
-            wallet.path,
-            "--wallet.name",
-            wallet.name,
-            "--wallet.hotkey",
-            "default",
-        )
-
-        yield
-
-        process.terminate()
+    def miner(self, wallet, netuid):
+        return self.Miner(self.dir, wallet, netuid)
 
-        await process.wait()
+    def validator(self, wallet, netuid):
+        return self.Validator(self.dir, wallet, netuid)