def test_agent_registration_and_id(monkeypatch, cloud_api): monkeypatch.setattr("prefect.agent.agent.Agent._verify_token", MagicMock()) monkeypatch.setattr("prefect.agent.agent.Client.register_agent", MagicMock(return_value="ID")) agent = Agent() assert agent._register_agent() == "ID" 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_catch_errors_in_heartbeat_thread(monkeypatch, cloud_api, caplog): """Check that errors in the heartbeat thread are caught, logged, and the thread keeps going""" monkeypatch.setattr( "prefect.agent.agent.Agent._submit_deploy_flow_run_jobs", MagicMock() ) monkeypatch.setattr( "prefect.agent.agent.Agent._setup_api_connection", MagicMock(return_value="id") ) heartbeat = MagicMock(side_effect=ValueError) monkeypatch.setattr("prefect.agent.agent.Agent.heartbeat", heartbeat) agent = Agent(max_polls=2) # Ignore registration agent._register_agent = MagicMock() agent.heartbeat_period = 0.1 agent.start() assert heartbeat.call_count > 1 assert any("Error in agent heartbeat" in m for m in caplog.messages)
def test_agent_poke_api(monkeypatch, cloud_api): import threading requests = pytest.importorskip("requests") def _poke_agent(agent_address): # May take a sec for the api server to startup for attempt in range(5): try: resp = requests.get(f"{agent_address}/api/health") break except Exception: time.sleep(0.1) else: assert False, "Failed to connect to health check" assert resp.status_code == 200 # Agent API is now available. Poke agent to start processing. requests.get(f"{agent_address}/api/poke") submit_deploy_flow_run_jobs = MagicMock() monkeypatch.setattr( "prefect.agent.agent.Agent._submit_deploy_flow_run_jobs", submit_deploy_flow_run_jobs, ) setup_api_connection = MagicMock(return_value="id") monkeypatch.setattr( "prefect.agent.agent.Agent._setup_api_connection", setup_api_connection ) heartbeat = MagicMock() monkeypatch.setattr("prefect.agent.agent.Agent.heartbeat", heartbeat) with socket.socket() as sock: sock.bind(("", 0)) port = sock.getsockname()[1] agent_address = f"http://127.0.0.1:{port}" # Poke agent in separate thread as main thread is blocked by main agent # process waiting for loop interval to complete. poke_agent_thread = threading.Thread(target=_poke_agent, args=(agent_address,)) poke_agent_thread.start() agent_start_time = time.time() agent = Agent(agent_address=agent_address, max_polls=1) # Ignore registration agent._register_agent = MagicMock() # Override loop interval to 5 seconds. agent._loop_intervals = {0: 5.0} agent.start() agent_stop_time = time.time() assert agent_stop_time - agent_start_time < 5.0 assert not agent._api_server_thread.is_alive() assert heartbeat.call_count == 1 assert submit_deploy_flow_run_jobs.call_count == 1 assert setup_api_connection.call_count == 1