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