feat(simulator): Add StopWhenSendLineMatch trigger

Author: msbrogli

hathor/simulator/trigger.py

@@ -16,6 +16,9 @@
 from typing import TYPE_CHECKING, Callable
 
 if TYPE_CHECKING:
+    from re import Match, Pattern
+
+    from hathor.p2p.protocol import HathorLineReceiver
     from hathor.simulator.fake_connection import FakeConnection
     from hathor.simulator.miner import AbstractMiner
     from hathor.simulator.tx_generator import RandomTransactionGenerator
@@ -107,3 +110,32 @@ def __init__(self, sub_triggers: list[Trigger]) -> None:
 
     def should_stop(self) -> bool:
         return all(trigger.should_stop() for trigger in self._sub_triggers)
+
+
+class StopWhenSendLineMatch(Trigger):
+    """Stop the simulation when the node sends a line that matches a designated regex pattern.
+    """
+
+    def __init__(self, protocol: 'HathorLineReceiver', regex: 'Pattern') -> None:
+        # patches protocol.sendLine
+        self.original_send_line = protocol.sendLine
+        setattr(protocol, 'sendLine', self._send_line_wrapper)
+
+        # regex pattern
+        self.regex = regex
+
+        # list of matches
+        self.matches: list['Match'] = []
+
+    def _send_line_wrapper(self, line: str) -> None:
+        """Check if line matches a designated regex pattern."""
+        self.original_send_line(line)
+        match = self.regex.match(line)
+        if match:
+            self.matches.append(match)
+
+    def should_stop(self) -> bool:
+        if self.matches:
+            self.matches = []
+            return True
+        return False

tests/simulation/test_trigger.py

@@ -1,5 +1,8 @@
-from hathor.simulator import Simulator
-from hathor.simulator.trigger import StopAfterMinimumBalance, StopAfterNMinedBlocks
+import re
+
+from hathor.p2p.messages import ProtocolMessages
+from hathor.simulator import FakeConnection, Simulator
+from hathor.simulator.trigger import StopAfterMinimumBalance, StopAfterNMinedBlocks, StopWhenSendLineMatch
 from tests import unittest
 
 
@@ -58,3 +61,13 @@ def test_stop_after_minimum_balance(self):
         self.assertLess(wallet.balance[token_uid].available, minimum_balance)
         self.assertTrue(self.simulator.run(3600, trigger=trigger))
         self.assertGreaterEqual(wallet.balance[token_uid].available, minimum_balance)
+
+    def test_stop_after_sendline(self):
+        manager2 = self.simulator.create_peer()
+        conn12 = FakeConnection(self.manager1, manager2, latency=0.05)
+        self.simulator.add_connection(conn12)
+
+        expected_prefix = f'^{ProtocolMessages.PEER_ID.value} '.encode('ascii')
+        regex = re.compile(expected_prefix)
+        trigger = StopWhenSendLineMatch(conn12._proto1, regex)
+        self.assertTrue(self.simulator.run(120, trigger=trigger))