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