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