refactor: use Hathor simulator on events simulator CLI

Author: glevco

hathor/cli/events_simulator/event_forwarding_websocket_factory.py

@@ -0,0 +1,32 @@
+#  Copyright 2023 Hathor Labs
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from typing import Any
+
+from twisted.internet.interfaces import IAddress
+
+from hathor.cli.events_simulator.event_forwarding_websocket_protocol import EventForwardingWebsocketProtocol
+from hathor.event.websocket import EventWebsocketFactory
+from hathor.simulator import Simulator
+
+
+class EventForwardingWebsocketFactory(EventWebsocketFactory):
+    def __init__(self, simulator: Simulator, *args: Any, **kwargs: Any) -> None:
+        self._simulator = simulator
+        super().__init__(*args, **kwargs)
+
+    def buildProtocol(self, _: IAddress) -> EventForwardingWebsocketProtocol:
+        protocol = EventForwardingWebsocketProtocol(self._simulator)
+        protocol.factory = self
+        return protocol

hathor/cli/events_simulator/event_forwarding_websocket_protocol.py

@@ -0,0 +1,47 @@
+#  Copyright 2023 Hathor Labs
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from typing import TYPE_CHECKING
+
+from autobahn.websocket import ConnectionRequest
+
+from hathor.event.websocket import EventWebsocketProtocol
+from hathor.simulator import Simulator
+
+if TYPE_CHECKING:
+    from hathor.cli.events_simulator.event_forwarding_websocket_factory import EventForwardingWebsocketFactory
+
+
+class EventForwardingWebsocketProtocol(EventWebsocketProtocol):
+    factory: 'EventForwardingWebsocketFactory'
+
+    def __init__(self, simulator: Simulator) -> None:
+        self._simulator = simulator
+        super().__init__()
+
+    def onConnect(self, request: ConnectionRequest) -> None:
+        super().onConnect(request)
+        self._simulator.run(60)
+
+    def onOpen(self) -> None:
+        super().onOpen()
+        self._simulator.run(60)
+
+    def onClose(self, wasClean: bool, code: int, reason: str) -> None:
+        super().onClose(wasClean, code, reason)
+        self._simulator.run(60)
+
+    def onMessage(self, payload: bytes, isBinary: bool) -> None:
+        super().onMessage(payload, isBinary)
+        self._simulator.run(60)

hathor/cli/events_simulator/events_simulator.py

@@ -12,10 +12,18 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import os
 from argparse import ArgumentParser, Namespace
 
+from autobahn.twisted.resource import WebSocketResource
+from structlog import get_logger
+from twisted.web.resource import Resource
+from twisted.web.server import Site
+
 DEFAULT_PORT = 8080
 
+logger = get_logger()
+
 
 def create_parser() -> ArgumentParser:
     from hathor.cli.events_simulator.scenario import Scenario
@@ -26,14 +34,17 @@ def create_parser() -> ArgumentParser:
 
     parser.add_argument('--scenario', help=f'One of {possible_scenarios}', type=str, required=True)
     parser.add_argument('--port', help='Port to run the WebSocket server', type=int, default=DEFAULT_PORT)
+    parser.add_argument('--seed', help='The seed used to create simulated events', type=int)
 
     return parser
 
 
 def execute(args: Namespace) -> None:
+    from hathor.conf import UNITTESTS_SETTINGS_FILEPATH
+    os.environ['HATHOR_CONFIG_YAML'] = UNITTESTS_SETTINGS_FILEPATH
+    from hathor.cli.events_simulator.event_forwarding_websocket_factory import EventForwardingWebsocketFactory
     from hathor.cli.events_simulator.scenario import Scenario
-    from hathor.event.storage import EventMemoryStorage
-    from hathor.event.websocket import EventWebsocketFactory
+    from hathor.simulator import Simulator
     from hathor.util import reactor
 
     try:
@@ -42,15 +53,36 @@ def execute(args: Namespace) -> None:
         possible_scenarios = [scenario.name for scenario in Scenario]
         raise ValueError(f'Invalid scenario "{args.scenario}". Choose one of {possible_scenarios}') from e
 
-    storage = EventMemoryStorage()
+    log = logger.new()
+    simulator = Simulator(args.seed)
+    simulator.start()
+    builder = simulator.get_default_builder() \
+        .disable_full_verification() \
+        .enable_event_queue()
+
+    manager = simulator.create_peer(builder)
+    event_ws_factory = manager._event_manager._event_ws_factory
+    assert event_ws_factory is not None
+
+    forwarding_ws_factory = EventForwardingWebsocketFactory(
+        simulator=simulator,
+        reactor=reactor,
+        event_storage=event_ws_factory._event_storage
+    )
+
+    manager._event_manager._event_ws_factory = forwarding_ws_factory
 
-    for event in scenario.value:
-        storage.save_event(event)
+    root = Resource()
+    api = Resource()
+    root.putChild(b'v1a', api)
+    api.putChild(b'event_ws', WebSocketResource(forwarding_ws_factory))
+    site = Site(root)
 
-    factory = EventWebsocketFactory(reactor, storage)
+    log.info('Started simulating events', scenario=args.scenario, seed=simulator.seed)
 
-    factory.start()
-    reactor.listenTCP(args.port, factory)
+    forwarding_ws_factory.start()
+    scenario.simulate(simulator, manager)
+    reactor.listenTCP(args.port, site)
     reactor.run()
 
 

hathor/cli/events_simulator/scenario.py

@@ -13,21 +13,83 @@
 #  limitations under the License.
 
 from enum import Enum
+from typing import TYPE_CHECKING
 
-from hathor.cli.events_simulator.scenarios.only_load_events import ONLY_LOAD_EVENTS
-from hathor.cli.events_simulator.scenarios.reorg_events import REORG_EVENTS
-from hathor.cli.events_simulator.scenarios.single_chain_blocks_and_transactions_events import (
-    SINGLE_CHAIN_BLOCKS_AND_TRANSACTIONS_EVENTS,
-)
-from hathor.cli.events_simulator.scenarios.single_chain_one_block_events import SINGLE_CHAIN_ONE_BLOCK_EVENTS
+if TYPE_CHECKING:
+    from hathor.manager import HathorManager
+    from hathor.simulator import Simulator
 
 
 class Scenario(Enum):
-    """
-    NOTE: The lists of events used in each scenario's enum value below were obtained from the tests in
-    tests.event.test_simulation.TestEventSimulation
-    """
-    ONLY_LOAD = ONLY_LOAD_EVENTS
-    SINGLE_CHAIN_ONE_BLOCK = SINGLE_CHAIN_ONE_BLOCK_EVENTS
-    SINGLE_CHAIN_BLOCKS_AND_TRANSACTIONS = SINGLE_CHAIN_BLOCKS_AND_TRANSACTIONS_EVENTS
-    REORG = REORG_EVENTS
+    ONLY_LOAD = 'ONLY_LOAD'
+    SINGLE_CHAIN_ONE_BLOCK = 'SINGLE_CHAIN_ONE_BLOCK'
+    SINGLE_CHAIN_BLOCKS_AND_TRANSACTIONS = 'SINGLE_CHAIN_BLOCKS_AND_TRANSACTIONS'
+    REORG = 'REORG'
+
+    def simulate(self, simulator: 'Simulator', manager: 'HathorManager') -> None:
+        simulate_fns = {
+            Scenario.ONLY_LOAD: simulate_only_load,
+            Scenario.SINGLE_CHAIN_ONE_BLOCK: simulate_single_chain_one_block,
+            Scenario.SINGLE_CHAIN_BLOCKS_AND_TRANSACTIONS: simulate_single_chain_blocks_and_transactions,
+            Scenario.REORG: simulate_reorg,
+        }
+
+        simulate_fn = simulate_fns[self]
+
+        simulate_fn(simulator, manager)
+
+
+def simulate_only_load(simulator: 'Simulator', _manager: 'HathorManager') -> None:
+    simulator.run(60)
+
+
+def simulate_single_chain_one_block(simulator: 'Simulator', manager: 'HathorManager') -> None:
+    from tests.utils import add_new_blocks
+    add_new_blocks(manager, 1)
+    simulator.run(60)
+
+
+def simulate_single_chain_blocks_and_transactions(simulator: 'Simulator', manager: 'HathorManager') -> None:
+    from hathor import daa
+    from hathor.conf import HathorSettings
+    from tests.utils import add_new_blocks, gen_new_tx
+
+    settings = HathorSettings()
+    assert manager.wallet is not None
+    address = manager.wallet.get_unused_address(mark_as_used=False)
+
+    add_new_blocks(manager, settings.REWARD_SPEND_MIN_BLOCKS + 1)
+    simulator.run(60)
+
+    tx = gen_new_tx(manager, address, 1000)
+    tx.weight = daa.minimum_tx_weight(tx)
+    tx.update_hash()
+    assert manager.propagate_tx(tx, fails_silently=False)
+    simulator.run(60)
+
+    tx = gen_new_tx(manager, address, 2000)
+    tx.weight = daa.minimum_tx_weight(tx)
+    tx.update_hash()
+    assert manager.propagate_tx(tx, fails_silently=False)
+    simulator.run(60)
+
+    add_new_blocks(manager, 1)
+    simulator.run(60)
+
+
+def simulate_reorg(simulator: 'Simulator', manager: 'HathorManager') -> None:
+    from hathor.simulator import FakeConnection
+    from tests.utils import add_new_blocks
+
+    builder = simulator.get_default_builder()
+    manager2 = simulator.create_peer(builder)
+
+    add_new_blocks(manager, 1)
+    simulator.run(60)
+
+    add_new_blocks(manager2, 2)
+    simulator.run(60)
+
+    connection = FakeConnection(manager, manager2)
+    simulator.add_connection(connection)
+    simulator.run(60)