Merge pull request #1147 from HathorNetwork/fix/quick_test-regression

fix(cli): subcommand quick_test not stopping after #1115

hathor/cli/quick_test.py

@@ -16,11 +16,49 @@
 from argparse import ArgumentParser
 from typing import Any
 
+from structlog import get_logger
+
 from hathor.cli.run_node import RunNode
 
+logger = get_logger()
 
-class QuickTest(RunNode):
 
+class VertexHandlerWrapper:
+    def __init__(self, vertex_handler, manager, n_blocks):
+        self.log = logger.new()
+        self._vertex_handler = vertex_handler
+        self._manager = manager
+        self._n_blocks = n_blocks
+
+    def on_new_vertex(self, *args: Any, **kwargs: Any) -> bool:
+        from hathor.transaction import Block
+        from hathor.transaction.base_transaction import GenericVertex
+
+        msg: str | None = None
+        res = self._vertex_handler.on_new_vertex(*args, **kwargs)
+
+        if self._n_blocks is None:
+            should_quit = res
+            msg = 'added a tx'
+        else:
+            vertex = args[0]
+            should_quit = False
+            assert isinstance(vertex, GenericVertex)
+
+            if isinstance(vertex, Block):
+                should_quit = vertex.get_height() >= self._n_blocks
+                msg = f'reached height {vertex.get_height()}'
+
+        if should_quit:
+            assert msg is not None
+            self.log.info(f'successfully {msg}, exit now')
+            self._manager.connections.disconnect_all_peers(force=True)
+            self._manager.reactor.fireSystemEvent('shutdown')
+            os._exit(0)
+        return res
+
+
+class QuickTest(RunNode):
     @classmethod
     def create_parser(cls) -> ArgumentParser:
         parser = super().create_parser()
@@ -30,38 +68,20 @@ def create_parser(cls) -> ArgumentParser:
         return parser
 
     def prepare(self, *, register_resources: bool = True) -> None:
-        from hathor.transaction import Block
-        from hathor.transaction.base_transaction import GenericVertex
+        from hathor.p2p.sync_v2.factory import SyncV2Factory
+        from hathor.p2p.sync_version import SyncVersion
+
         super().prepare(register_resources=False)
         self._no_wait = self._args.no_wait
 
-        self.log.info('patching on_new_tx to quit on success')
-        orig_on_new_tx = self.manager.on_new_tx
-
-        def patched_on_new_tx(*args: Any, **kwargs: Any) -> bool:
-            res = orig_on_new_tx(*args, **kwargs)
-            msg: str | None = None
-
-            if self._args.quit_after_n_blocks is None:
-                should_quit = res
-                msg = 'added a tx'
-            else:
-                vertex = args[0]
-                should_quit = False
-                assert isinstance(vertex, GenericVertex)
-
-                if isinstance(vertex, Block):
-                    should_quit = vertex.get_height() >= self._args.quit_after_n_blocks
-                    msg = f'reached height {vertex.get_height()}'
-
-            if should_quit:
-                assert msg is not None
-                self.log.info(f'successfully {msg}, exit now')
-                self.manager.connections.disconnect_all_peers(force=True)
-                self.reactor.fireSystemEvent('shutdown')
-                os._exit(0)
-            return res
-        self.manager.on_new_tx = patched_on_new_tx
+        self.log.info('patching vertex_handler.on_new_vertex to quit on success')
+        p2p_factory = self.manager.connections.get_sync_factory(SyncVersion.V2)
+        assert isinstance(p2p_factory, SyncV2Factory)
+        p2p_factory.vertex_handler = VertexHandlerWrapper(
+            self.manager.vertex_handler,
+            self.manager,
+            self._args.quit_after_n_blocks,
+        )  # type: ignore
 
         timeout = 300
         self.log.info('exit with error code if it take too long', timeout=timeout)