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 token check and test graphql query
agent._verify_token = MagicMock()
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 the token check and registration
agent._verify_token = MagicMock()
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()
