def workflow_instance( test_workflow: Workflow, ) -> Tuple[RunningTask, Subscriber, Subscriber, Dict[str, List[Subscriber]]]: td = timedelta(seconds=100) sub1 = Subscription("start_collect", True, td) sub2 = Subscription("collect", True, td) sub3 = Subscription("collect_done", True, td) s1 = Subscriber.from_list("s1", [sub1, sub2, sub3]) s2 = Subscriber.from_list("s2", [sub2, sub3]) subscriptions = { "start_collect": [s1], "collect": [s1, s2], "collect_done": [s1, s2] } w, _ = RunningTask.empty(test_workflow, lambda: subscriptions) w.received_messages = [ Action("start_collect", w.id, "start"), ActionDone("start_collect", w.id, "start", s1.id), ActionDone("start_collect", w.id, "start", s2.id), Event("godot", { "a": 1, "b": 2 }), Action("collect", w.id, "collect"), ActionDone("collect", w.id, "collect", s1.id), ] w.move_to_next_state() return w, s1, s2, subscriptions
def test_message_serialization() -> None: roundtrip(Event("test", {"a": "b", "c": 1, "d": "bla"})) roundtrip(Action("test", "123", "step_name")) roundtrip(Action("test", "123", "step_name", {"test": 1})) roundtrip(ActionDone("test", "123", "step_name", "sub")) roundtrip(ActionDone("test", "123", "step_name", "sub", {"test": 1})) roundtrip(ActionError("test", "123", "step_name", "sub", "oops")) roundtrip( ActionError("test", "123", "step_name", "sub", "oops", {"test": 23}))
def test_message_serialization() -> None: task_id = TaskId("123") subsctiber_id = SubscriberId("sub") roundtrip(Event("test", {"a": "b", "c": 1, "d": "bla"})) roundtrip(Action("test", task_id, "step_name")) roundtrip(Action("test", task_id, "step_name", {"test": 1})) roundtrip(ActionDone("test", task_id, "step_name", subsctiber_id)) roundtrip( ActionDone("test", task_id, "step_name", subsctiber_id, {"test": 1})) roundtrip(ActionError("test", task_id, "step_name", subsctiber_id, "oops")) roundtrip( ActionError("test", task_id, "step_name", subsctiber_id, "oops", {"test": 23}))
async def test_handler_invocation( merge_handler: MergeOuterEdgesHandler, subscription_handler: SubscriptionHandler, message_bus: MessageBus, ) -> None: merge_called: asyncio.Future[TaskId] = asyncio.get_event_loop( ).create_future() def mocked_merge(task_id: TaskId) -> None: merge_called.set_result(task_id) # monkey patching the merge_outer_edges method # use setattr here, since assignment does not work in mypy https://github.com/python/mypy/issues/2427 setattr(merge_handler, "merge_outer_edges", mocked_merge) subscribers = await subscription_handler.list_subscriber_for( merge_outer_edges) assert subscribers[0].id == "resotocore" task_id = TaskId("test_task_1") await message_bus.emit( Action(merge_outer_edges, task_id, merge_outer_edges)) assert await merge_called == task_id
def test_pending_action_for( workflow_instance: Tuple[RunningTask, Subscriber, Subscriber, Dict[str, List[Subscriber]]], ) -> None: wi, s1, s2, subscriptions = workflow_instance # s1 already sent a done message for the current step assert wi.pending_action_for(s1) is None # s2 is still expected to provide a done message assert wi.pending_action_for(s2) == Action("collect", wi.id, "collect")
def commands_to_execute(self) -> Sequence[TaskCommand]: """ When the state is entered, emit the action message and inform all actors. """ return [ SendMessage( Action(self.perform.message_type, self.instance.id, self.step.name)) ]
def pending_action_for(self, subscriber: Subscriber) -> Optional[Action]: """ In case this task is waiting for an action result from the given subscriber, the relevant action is returned. """ state = self.current_state if isinstance(state, PerformActionState): message_type = state.perform.message_type subscriptions = state.wait_for if subscriber in subscriptions and self.ack_for( message_type, subscriber) is None: return Action(message_type, self.id, state.step.name) return None
async def test_recover_workflow( running_task_db: RunningTaskDb, job_db: JobDb, message_bus: MessageBus, event_sender: AnalyticsEventSender, subscription_handler: SubscriptionHandler, all_events: List[Message], cli: CLI, test_workflow: Workflow, ) -> None: def handler() -> TaskHandlerService: th = TaskHandlerService( running_task_db, job_db, message_bus, event_sender, subscription_handler, Scheduler(), cli, empty_config(), ) th.task_descriptions = [test_workflow] return th await subscription_handler.add_subscription(SubscriberId("sub_1"), "start_collect", True, timedelta(seconds=30)) sub1 = await subscription_handler.add_subscription(SubscriberId("sub_1"), "collect", True, timedelta(seconds=30)) sub2 = await subscription_handler.add_subscription(SubscriberId("sub_2"), "collect", True, timedelta(seconds=30)) async with handler() as wf1: # kick off a new workflow await wf1.handle_event(Event("start me up")) assert len(wf1.tasks) == 1 # expect a start_collect action message a: Action = await wait_for_message(all_events, "start_collect", Action) await wf1.handle_action_done(ActionDone(a.message_type, a.task_id, a.step_name, sub1.id, dict(a.data))) # expect a collect action message b: Action = await wait_for_message(all_events, "collect", Action) await wf1.handle_action_done(ActionDone(b.message_type, b.task_id, b.step_name, sub1.id, dict(b.data))) # subscriber 3 is also registering for collect # since the collect phase is already started, it should not participate in this round sub3 = await subscription_handler.add_subscription(SubscriberId("sub_3"), "collect", True, timedelta(seconds=30)) # simulate a restart, wf1 is stopped and wf2 needs to recover from database async with handler() as wf2: assert len(wf2.tasks) == 1 wfi = list(wf2.tasks.values())[0] assert wfi.current_state.name == "act" assert (await wf2.list_all_pending_actions_for(sub1)) == [] assert (await wf2.list_all_pending_actions_for(sub2)) == [Action("collect", wfi.id, "act", {})] assert (await wf2.list_all_pending_actions_for(sub3)) == [] await wf2.handle_action_done(ActionDone("collect", wfi.id, "act", sub2.id, {})) # expect an event workflow_end await wait_for_message(all_events, "task_end", Event) # all workflow instances are gone assert len(wf2.tasks) == 0 # simulate a restart, wf3 should start from a clean slate, since all instances are done async with handler() as wf3: assert len(wf3.tasks) == 0