def test_setup_api_connection_attaches_agent_id(cloud_api):
agent = Agent(max_polls=1)
# Return a fake id from the "backend"
agent.client.register_agent = MagicMock(return_value="ID")
# Ignore the test graphql query
agent.client.graphql = MagicMock()
agent._setup_api_connection()
assert agent.client._attached_headers == {"X-PREFECT-AGENT-ID": "ID"}
def test_setup_api_connection_runs_test_query(test_query_succeeds, cloud_api):
agent = Agent()
# Ignore registration
agent._register_agent = MagicMock()
if test_query_succeeds:
# Create a successful test query
agent.client.graphql = MagicMock(return_value="Hello")
with nullcontext() if test_query_succeeds else pytest.raises(Exception):
agent._setup_api_connection()
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
