diff --git a/hathor/cli/quick_test.py b/hathor/cli/quick_test.py index 1ba6fd0ff..2bf6f16fe 100644 --- a/hathor/cli/quick_test.py +++ b/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)