def test_mark_flow_as_submitted(monkeypatch, cloud_api, with_task_runs):
agent = Agent()
agent.client = MagicMock()
agent._mark_flow_as_submitted(flow_run=GraphQLResult({
"id":
"id",
"serialized_state":
Scheduled().serialize(),
"version":
1,
"task_runs": ([
GraphQLResult({
"id": "task-id",
"version": 1,
"serialized_state": Scheduled().serialize(),
})
] if with_task_runs else []),
}))
agent.client.set_flow_run_state.assert_called_once_with(
flow_run_id="id",
version=1,
state=Submitted(message="Submitted for execution"))
if with_task_runs:
agent.client.set_task_run_state.assert_called_once_with(
task_run_id="task-id",
version=1,
state=Submitted(message="Submitted for execution"),
)
else:
agent.client.set_task_run_state.assert_not_called()
def test_mark_flow_as_failed(monkeypatch, cloud_api):
agent = Agent()
agent.client = MagicMock()
agent._mark_flow_as_failed(
flow_run=GraphQLResult({
"id": "id",
"serialized_state": Scheduled().serialize(),
"version": 1,
"task_runs": [],
}),
message="foo",
)
agent.client.set_flow_run_state.assert_called_with(
flow_run_id="id", version=1, state=Failed(message="foo"))
