def test_agent_poke_api(monkeypatch, runner_token, 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)
# 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
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(max_polls=1)
agent.start()
assert agent._register_agent() == "ID"
assert agent.client._attached_headers == {"X-PREFECT-AGENT-ID": "ID"}
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
def test_agent_rerieve_config(monkeypatch, cloud_api):
monkeypatch.setattr("prefect.agent.agent.Agent._verify_token", MagicMock())
monkeypatch.setattr("prefect.agent.agent.Client.register_agent",
MagicMock(return_value="ID"))
monkeypatch.setattr(
"prefect.agent.agent.Client.get_agent_config",
MagicMock(return_value={"settings": "yes"}),
)
agent = Agent(max_polls=1)
agent.start()
assert agent._retrieve_agent_config() == {"settings": "yes"}
def test_catch_errors_in_heartbeat_thread(monkeypatch, runner_token, cloud_api, caplog):
"""Check that errors in the heartbeat thread are caught, logged, and the thread keeps going"""
monkeypatch.setattr("prefect.agent.agent.Agent.agent_process", MagicMock())
monkeypatch.setattr(
"prefect.agent.agent.Agent.agent_connect", MagicMock(return_value="id")
)
heartbeat = MagicMock(side_effect=ValueError)
monkeypatch.setattr("prefect.agent.agent.Agent.heartbeat", heartbeat)
agent = Agent(max_polls=2)
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_start_max_polls_zero(monkeypatch, runner_token, cloud_api):
on_shutdown = MagicMock()
monkeypatch.setattr("prefect.agent.agent.Agent.on_shutdown", on_shutdown)
agent_process = MagicMock()
monkeypatch.setattr("prefect.agent.agent.Agent.agent_process", agent_process)
agent_connect = MagicMock(return_value="id")
monkeypatch.setattr("prefect.agent.agent.Agent.agent_connect", agent_connect)
heartbeat = MagicMock()
monkeypatch.setattr("prefect.agent.agent.Agent.heartbeat", heartbeat)
agent = Agent(max_polls=0)
agent.start()
assert on_shutdown.call_count == 1
assert agent_process.call_count == 0
assert heartbeat.call_count == 0
def test_agent_start_max_polls(monkeypatch, runner_token):
on_shutdown = MagicMock()
monkeypatch.setattr("prefect.agent.agent.Agent.on_shutdown", on_shutdown)
agent_process = MagicMock()
monkeypatch.setattr("prefect.agent.agent.Agent.agent_process", agent_process)
agent_connect = MagicMock(return_value="id")
monkeypatch.setattr("prefect.agent.agent.Agent.agent_connect", agent_connect)
heartbeat = MagicMock()
monkeypatch.setattr("prefect.agent.agent.Agent.heartbeat", heartbeat)
agent = Agent(max_polls=1)
agent.start()
assert on_shutdown.called
assert agent_process.called
assert agent_process.call_args[0][1] == "id"
assert heartbeat.called
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)
