diff --git a/hathor/cli/events_simulator/events_simulator.py b/hathor/cli/events_simulator/events_simulator.py index 7d09f1cfa..d4beea9de 100644 --- a/hathor/cli/events_simulator/events_simulator.py +++ b/hathor/cli/events_simulator/events_simulator.py @@ -22,22 +22,29 @@ def create_parser() -> ArgumentParser: from hathor.cli.util import create_parser parser = create_parser() - possible_scenarios = [scenario.value for scenario in Scenario] + possible_scenarios = [scenario.name for scenario in Scenario] - parser.add_argument('--scenario', help=f'One of {possible_scenarios}', type=Scenario, required=True) + 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) return parser def execute(args: Namespace) -> None: + from hathor.cli.events_simulator.scenario import Scenario from hathor.event.storage import EventMemoryStorage from hathor.event.websocket import EventWebsocketFactory from hathor.util import reactor + try: + scenario = Scenario[args.scenario] + except KeyError as e: + possible_scenarios = [scenario.name for scenario in Scenario] + raise ValueError(f'Invalid scenario "{args.scenario}". Choose one of {possible_scenarios}') from e + storage = EventMemoryStorage() - for event in args.scenario.value: + for event in scenario.value: storage.save_event(event) factory = EventWebsocketFactory(reactor, storage)