diff --git a/hathor/builder/builder.py b/hathor/builder/builder.py index c074ecf2f..e298993cb 100644 --- a/hathor/builder/builder.py +++ b/hathor/builder/builder.py @@ -500,6 +500,16 @@ def set_full_verification(self, full_verification: bool) -> 'Builder': self._full_verification = full_verification return self + def enable_full_verification(self) -> 'Builder': + self.check_if_can_modify() + self._full_verification = True + return self + + def disable_full_verification(self) -> 'Builder': + self.check_if_can_modify() + self._full_verification = False + return self + def set_soft_voided_tx_ids(self, soft_voided_tx_ids: Set[bytes]) -> 'Builder': self.check_if_can_modify() self._soft_voided_tx_ids = soft_voided_tx_ids diff --git a/hathor/simulator/simulator.py b/hathor/simulator/simulator.py index 8320d6c78..a4d60744b 100644 --- a/hathor/simulator/simulator.py +++ b/hathor/simulator/simulator.py @@ -15,17 +15,15 @@ import secrets import time from collections import OrderedDict -from typing import TYPE_CHECKING, Any, Generator, List, Optional, Set +from typing import TYPE_CHECKING, Any, Generator, List, Optional from mnemonic import Mnemonic from structlog import get_logger -from hathor.builder import Builder +from hathor.builder import BuildArtifacts, Builder from hathor.conf import HathorSettings from hathor.daa import TestMode, _set_test_mode -from hathor.event.websocket import EventWebsocketFactory from hathor.manager import HathorManager -from hathor.p2p.peer_id import PeerId from hathor.simulator.clock import HeapClock from hathor.simulator.miner.dummy_miner import DummyMiner from hathor.simulator.miner.geometric_miner import GeometricMiner @@ -141,38 +139,41 @@ def stop(self) -> None: self._started = False self._patches_rc_decrement() - def create_peer( - self, - network: Optional[str] = None, - peer_id: Optional[PeerId] = None, - enable_sync_v1: bool = True, - enable_sync_v2: bool = True, - soft_voided_tx_ids: Optional[Set[bytes]] = None, - full_verification: bool = True, - event_ws_factory: Optional[EventWebsocketFactory] = None - ) -> HathorManager: + def get_default_builder(self) -> Builder: + """ + Returns a builder with default configuration, for convenience when using create_peer() or create_artifacts() + """ + return Builder() \ + .set_network(self._network) \ + .set_soft_voided_tx_ids(set()) \ + .enable_full_verification() \ + .enable_sync_v1() \ + .enable_sync_v2() \ + .use_memory() + + def create_peer(self, builder: Builder) -> HathorManager: + """ + Returns a manager from a builder, after configuring it for simulator use. + You may get a builder from get_default_builder() for convenience. + """ + artifacts = self.create_artifacts(builder) + return artifacts.manager + + def create_artifacts(self, builder: Builder) -> BuildArtifacts: + """ + Returns build artifacts from a builder, after configuring it for simulator use. + You may get a builder from get_default_builder() for convenience. + """ assert self._started, 'Simulator is not started.' - assert peer_id is not None # XXX: temporary, for checking that tests are using the peer_id wallet = HDWallet(gap_limit=2) wallet._manually_initialize() - builder = Builder() \ + artifacts = builder \ .set_reactor(self._clock) \ - .set_peer_id(peer_id or PeerId()) \ - .set_network(network or self._network) \ - .set_wallet(wallet) \ .set_rng(Random(self.rng.getrandbits(64))) \ - .set_enable_sync_v1(enable_sync_v1) \ - .set_enable_sync_v2(enable_sync_v2) \ - .set_full_verification(full_verification) \ - .set_soft_voided_tx_ids(soft_voided_tx_ids or set()) \ - .use_memory() - - if event_ws_factory: - builder.enable_event_manager(event_ws_factory=event_ws_factory) - - artifacts = builder.build() + .set_wallet(wallet) \ + .build() artifacts.manager.start() self.run_to_completion() @@ -184,7 +185,7 @@ def create_peer( self.log.debug('randomized step: generate wallet', words=words) wallet.unlock(words=words, tx_storage=artifacts.tx_storage) - return artifacts.manager + return artifacts def create_tx_generator(self, peer: HathorManager, *args: Any, **kwargs: Any) -> RandomTransactionGenerator: return RandomTransactionGenerator(peer, self.rng, *args, **kwargs) diff --git a/tests/event/test_simulation.py b/tests/event/test_simulation.py index 0e736dcb8..57b4742cd 100644 --- a/tests/event/test_simulation.py +++ b/tests/event/test_simulation.py @@ -38,11 +38,11 @@ def test_only_load(self): simulator.start() main_peer_id = PeerId() - main_manager = simulator.create_peer( - peer_id=main_peer_id, - full_verification=False, - event_ws_factory=Mock() - ) + builder = simulator.get_default_builder() \ + .set_peer_id(main_peer_id) \ + .disable_full_verification() \ + .enable_event_manager(event_ws_factory=Mock()) + main_manager = simulator.create_peer(builder) main_manager.allow_mining_without_peers() simulator.run(2 * 60) @@ -102,11 +102,11 @@ def test_single_chain_one_block(self): simulator.start() main_peer_id = PeerId() - main_manager = simulator.create_peer( - peer_id=main_peer_id, - full_verification=False, - event_ws_factory=Mock() - ) + builder = simulator.get_default_builder() \ + .set_peer_id(main_peer_id) \ + .disable_full_verification() \ + .enable_event_manager(event_ws_factory=Mock()) + main_manager = simulator.create_peer(builder) main_manager.allow_mining_without_peers() miner = simulator.create_dummy_miner(main_manager, block_times=[120]) @@ -202,11 +202,11 @@ def test_single_chain_blocks_and_transactions(self): simulator.start() main_peer_id = PeerId() - main_manager = simulator.create_peer( - peer_id=main_peer_id, - full_verification=False, - event_ws_factory=Mock() - ) + builder = simulator.get_default_builder() \ + .set_peer_id(main_peer_id) \ + .disable_full_verification() \ + .enable_event_manager(event_ws_factory=Mock()) + main_manager = simulator.create_peer(builder) main_manager.allow_mining_without_peers() miner = simulator.create_dummy_miner(main_manager, block_times=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 120]) @@ -303,11 +303,11 @@ def test_reorg(self): simulator.start() main_peer_id = PeerId() - main_manager = simulator.create_peer( - peer_id=main_peer_id, - full_verification=False, - event_ws_factory=Mock() - ) + builder = simulator.get_default_builder() \ + .set_peer_id(main_peer_id) \ + .disable_full_verification() \ + .enable_event_manager(event_ws_factory=Mock()) + main_manager = simulator.create_peer(builder) main_manager.allow_mining_without_peers() miner = simulator.create_dummy_miner(main_manager, block_times=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 120]) diff --git a/tests/simulation/base.py b/tests/simulation/base.py index 91a4a1f2b..8acb087ca 100644 --- a/tests/simulation/base.py +++ b/tests/simulation/base.py @@ -35,9 +35,11 @@ def create_peer(self, enable_sync_v1=None, enable_sync_v2=None, soft_voided_tx_i assert enable_sync_v1 or enable_sync_v2, 'enable at least one sync version' if simulator is None: simulator = self.simulator - return simulator.create_peer( - peer_id=self.get_random_peer_id_from_pool(rng=simulator.rng), - soft_voided_tx_ids=soft_voided_tx_ids, - enable_sync_v1=enable_sync_v1, - enable_sync_v2=enable_sync_v2, - ) + + builder = simulator.get_default_builder() \ + .set_peer_id(self.get_random_peer_id_from_pool(rng=simulator.rng)) \ + .set_soft_voided_tx_ids(soft_voided_tx_ids) \ + .set_enable_sync_v1(enable_sync_v1) \ + .set_enable_sync_v2(enable_sync_v2) + + return simulator.create_peer(builder) diff --git a/tests/simulation/test_simulator_itself.py b/tests/simulation/test_simulator_itself.py index 00b6b10d0..6683a37b4 100644 --- a/tests/simulation/test_simulator_itself.py +++ b/tests/simulation/test_simulator_itself.py @@ -46,11 +46,13 @@ def create_simulator_peer(self, simulator, peer_id_pool, enable_sync_v1=None, en 'the test class or pass `enable_sync_v2` by argument') enable_sync_v2 = self._enable_sync_v2 assert enable_sync_v1 or enable_sync_v2, 'enable at least one sync version' - return simulator.create_peer( - peer_id=self.get_random_peer_id_from_pool(), - enable_sync_v1=enable_sync_v1, - enable_sync_v2=enable_sync_v2, - ) + + builder = simulator.get_default_builder() \ + .set_peer_id(self.get_random_peer_id_from_pool()) \ + .set_enable_sync_v1(enable_sync_v1) \ + .set_enable_sync_v2(enable_sync_v2) + + return simulator.create_peer(builder) def _simulate_run(self, run_i, simulator): # XXX: the following was adapted from test_new_syncing_peer, it doesn't matter too much, but has good coverage diff --git a/tests/simulation/test_trigger.py b/tests/simulation/test_trigger.py index 572d1a319..728b04a96 100644 --- a/tests/simulation/test_trigger.py +++ b/tests/simulation/test_trigger.py @@ -12,7 +12,8 @@ def setUp(self): self.simulator.start() peer_id = PeerId() - self.manager1 = self.simulator.create_peer(peer_id=peer_id) + builder = self.simulator.get_default_builder().set_peer_id(peer_id) + self.manager1 = self.simulator.create_peer(builder) self.manager1.allow_mining_without_peers() print('-' * 30)