def get_task_run_info(self, flow_run_id, task_id, map_index, *args,
**kwargs):
"""
Return task run if found, otherwise create it
"""
self.call_count["get_task_run_info"] += 1
task_run = next(
(t for t in self.task_runs.values() if t.flow_run_id == flow_run_id
and t.task_id == task_id and t.map_index == map_index),
None,
)
if not task_run:
task_run = TaskRun(
id=str(uuid.uuid4()),
task_id=task_id,
flow_run_id=flow_run_id,
map_index=map_index,
)
self.task_runs[task_run.id] = task_run
return TaskRunInfoResult(
id=task_run.id,
task_id=task_id,
version=task_run.version,
state=task_run.state,
)
def get_flow_run_info(self, flow_run_id, *args, **kwargs):
self.call_count["get_flow_run_info"] += 1
flow_run = self.flow_runs[flow_run_id]
task_runs = [t for t in self.task_runs.values() if t.flow_run_id == flow_run_id]
return FlowRunInfoResult(
id=flow_run.id,
flow_id=flow_run.flow_id,
parameters={},
context=None,
version=flow_run.version,
scheduled_start_time=pendulum.parse("2019-01-25T19:15:58.632412+00:00"),
state=flow_run.state,
task_runs=[
TaskRunInfoResult(
id=tr.id,
task_id=tr.task_slug,
task_slug=tr.task_slug,
version=tr.version,
state=tr.state,
)
for tr in task_runs
],
)