def test_submit_deploy_flow_run_jobs_raises_exception_and_logs(monkeypatch, cloud_api):
client = MagicMock()
client.return_value.graphql.side_effect = ValueError("Error")
monkeypatch.setattr("prefect.agent.agent.Client", client)
executor = MagicMock()
agent = Agent()
with pytest.raises(Exception):
agent._submit_deploy_flow_run_jobs(executor, "id")
assert client.write_run_log.called
def test_agent_start_max_polls(cloud_api, max_polls):
agent = Agent(max_polls=max_polls)
# Mock the backend API to avoid immediate failure
agent._setup_api_connection = MagicMock(return_value="id")
# Mock the deployment func to count calls
agent._submit_deploy_flow_run_jobs = MagicMock()
agent.start()
agent._submit_deploy_flow_run_jobs.call_count == max_polls
def test_submit_deploy_flow_run_jobs_no_runs_found(monkeypatch, cloud_api):
gql_return = MagicMock(return_value=MagicMock(data=MagicMock(
set_flow_run_state=None,
set_task_run_state=None,
get_runs_in_queue=MagicMock(flow_run_ids=["id"]),
flow_run=[],
)))
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.agent.agent.Client", client)
executor = MagicMock()
agent = Agent()
assert not agent._submit_deploy_flow_run_jobs(executor)
assert not executor.submit.called
def test_submit_deploy_flow_run_jobs(monkeypatch, cloud_api):
gql_return = MagicMock(
return_value=MagicMock(
data=MagicMock(
set_flow_run_state=None,
set_task_run_state=None,
get_runs_in_queue=MagicMock(flow_run_ids=["id"]),
flow_run=[
GraphQLResult(
{
"id": "id",
"serialized_state": Scheduled().serialize(),
"version": 1,
"task_runs": [
GraphQLResult(
{
"id": "id",
"version": 1,
"serialized_state": Scheduled().serialize(),
}
)
],
"scheduled_start_time": str(pendulum.now()),
}
)
],
)
)
)
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.agent.agent.Client", client)
executor = MagicMock()
future_mock = MagicMock()
executor.submit = MagicMock(return_value=future_mock)
agent = Agent()
assert agent._submit_deploy_flow_run_jobs(executor)
assert executor.submit.called
assert future_mock.add_done_callback.called
