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"))