async def test_wait_for_running_job(
task_handler: TaskHandlerService, test_workflow: Workflow, all_events: List[Message]
) -> None:
test_workflow.on_surpass = TaskSurpassBehaviour.Wait
task_handler.task_descriptions = [test_workflow]
# subscribe as collect handler - the workflow will need to wait for this handler
sub = await task_handler.subscription_handler.add_subscription(
SubscriberId("sub_1"), "collect", True, timedelta(seconds=30)
)
await task_handler.handle_event(Event("start me up"))
# check, that the workflow has started
running_before = await task_handler.running_tasks()
assert len(running_before) == 1
act: Action = await wait_for_message(all_events, "collect", Action)
# pull the same trigger: the workflow can not be started, since there is already one in progress -> wait
await task_handler.handle_event(Event("start me up"))
# report success of the only subscriber
await task_handler.handle_action_done(ActionDone("collect", act.task_id, act.step_name, sub.id, dict(act.data)))
# check overdue tasks: wipe finished tasks and eventually start waiting tasks
await task_handler.check_overdue_tasks()
# check, that the workflow has started
running_after = await task_handler.running_tasks()
assert len(running_after) == 1
t_before, t_after = running_before[0], running_after[0]
assert t_before.descriptor.id == t_after.descriptor.id and t_before.id != t_after.id
def test_complete_workflow(
workflow_instance: Tuple[RunningTask, Subscriber, Subscriber, Dict[str, List[Subscriber]]]
) -> None:
init, s1, s2, subscriptions = workflow_instance
# start new workflow instance
wi, events = RunningTask.empty(init.descriptor, lambda: subscriptions)
assert wi.current_step.name == "start"
assert len(events) == 2
events = wi.handle_done(ActionDone("start", wi.id, "start", s1.id))
assert wi.current_step.name == "wait"
assert len(events) == 0
handled, events = wi.handle_event(Event("godot", {"a": 2}))
assert wi.current_step.name == "wait"
assert handled is False
assert len(events) == 0
handled, events = wi.handle_event(Event("godot", {"a": 1, "d": "test"}))
assert wi.current_step.name == "collect"
assert handled is True
assert len(events) == 2 # event from EmitEvent and action from PerformAction
events = wi.handle_done(ActionDone("start", wi.id, "start", s1.id)) #
assert wi.current_step.name == "collect"
assert len(events) == 0
events = wi.handle_done(ActionDone("collect", wi.id, "collect", s1.id))
assert wi.current_step.name == "collect"
assert len(events) == 0
events = wi.handle_done(ActionDone("collect", wi.id, "collect", s2.id))
assert wi.current_step.name == "done"
assert len(events) == 1
events = wi.handle_done(ActionDone("done", wi.id, "done", s1.id))
assert len(events) == 0
assert wi.current_step.name == "done"
events = wi.handle_done(ActionDone("done", wi.id, "done", s2.id))
assert len(events) == 1
assert wi.is_active is False
def test_eq() -> None:
s1 = Step("a", PerformAction("a"), timedelta())
s2 = Step("a", WaitForEvent("a", {"foo": "bla"}), timedelta())
s3 = Step("a", EmitEvent(Event("a", {"a": "b"})), timedelta())
s4 = Step("a", ExecuteCommand("echo hello"), timedelta())
assert s1 == Step("a", PerformAction("a"), timedelta())
assert s2 == Step("a", WaitForEvent("a", {"foo": "bla"}), timedelta())
assert s3 == Step("a", EmitEvent(Event("a", {"a": "b"})), timedelta())
assert s4 == Step("a", ExecuteCommand("echo hello"), timedelta())
trigger = [EventTrigger("start me up")]
assert Workflow("a", "a", [s1, s2, s3, s4],
trigger) == Workflow("a", "a", [s1, s2, s3, s4], trigger)
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}))
def test_workflow() -> Workflow:
return Workflow(
TaskDescriptorId("test_workflow"),
"Speakable name of workflow",
[
Step("start", PerformAction("start_collect"), timedelta(seconds=10)),
Step("wait", WaitForEvent("godot", {"a": 1}), timedelta(seconds=10)),
Step("emit_event", EmitEvent(Event("hello", {"a": 1})), timedelta(seconds=10)),
Step("collect", PerformAction("collect"), timedelta(seconds=10)),
Step("done", PerformAction("collect_done"), timedelta(seconds=10), StepErrorBehaviour.Stop),
],
[EventTrigger("start me up")],
)
def empty(
descriptor: TaskDescription,
subscriber_by_event: Callable[[], Dict[str, List[Subscriber]]]
) -> Tuple[RunningTask, Sequence[TaskCommand]]:
assert len(
descriptor.steps) > 0, "TaskDescription needs at least one step!"
uid = str(uuid.uuid1())
wi = RunningTask(uid, descriptor, subscriber_by_event)
messages = [
SendMessage(Event("task_started", data={"task": descriptor.name})),
*wi.move_to_next_state()
]
return wi, messages
async def emit() -> None:
await message_bus.emit(Event("foo"))
await message_bus.emit(Event("foo"))
await message_bus.emit(Event("bla"))
await message_bus.emit(Event("bar"))
def test_marshalling_task_command() -> None:
roundtrip(SendMessage(Event("test", {"foo": "hello"})))
roundtrip(ExecuteOnCLI("test", frozendict({"fii": "bla"})))
def test_marshalling_step_action() -> None:
roundtrip(PerformAction("test"))
roundtrip(EmitEvent(Event("test", {"foo": "hello"})))
roundtrip(WaitForEvent("test", {"foo": "hello"}))
roundtrip(ExecuteCommand("help"))
def __init__(self, instance: RunningTask):
self.event = Event("task_end")
super().__init__(Step("task_end", EmitEvent(self.event)), instance)
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
async def test_run_job(task_handler: TaskHandlerService, all_events: List[Message]) -> None:
await task_handler.handle_event(Event("start me up"))
started: Event = await wait_for_message(all_events, "task_started", Event)
await wait_for_message(all_events, "task_end", Event)
assert started.data["task"] == "Speakable name of workflow"
