def test_deploy_flows_passes_base_agent(runner_token, cloud_api):
agent = Agent()
with pytest.raises(NotImplementedError):
agent.deploy_flow(None)
def test_agent_log_level_debug(cloud_api):
with set_temporary_config(
{"cloud.agent.auth_token": "TEST_TOKEN", "cloud.agent.level": "DEBUG"}
):
agent = Agent()
assert agent.logger.level == 10
def test_multiple_agent_init_doesnt_duplicate_logs(runner_token, cloud_api):
a, b, c = Agent(), Agent(), Agent()
assert len(c.logger.handlers) == 1
def test_agent_labels(cloud_api):
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = Agent(labels=["test", "2"])
assert agent.labels == ["test", "2"]
def test_agent_fails_no_auth_token(cloud_api):
with pytest.raises(RuntimeError,
match="Error while contacting API") as err:
Agent().start()
assert isinstance(err.value.__cause__, AuthorizationError)
def test_agent_log_level(runner_token, cloud_api):
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = Agent()
assert agent.logger.level == 20
def test_agent_labels_from_config_var(runner_token):
with set_temporary_config({"cloud.agent.labels": "['test', '2']"}):
agent = Agent()
assert agent.labels == ["test", "2"]
def test_agent_retrieve_config_requires_config_id_set(cloud_api):
agent = Agent(max_polls=1)
with pytest.raises(ValueError, match="agent_config_id"):
assert agent._retrieve_agent_config() == {"settings": "yes"}
def test_agent_init(runner_token):
agent = Agent()
assert agent
def test_heartbeat_is_noop_by_default(cloud_api):
agent = Agent()
assert not agent.heartbeat()
def test_deploy_flow_run_completed_callback_removes_id_from_submitted(
cloud_api):
agent = Agent()
agent.submitting_flow_runs.add("id")
agent._deploy_flow_run_completed_callback(None, "id")
assert len(agent.submitting_flow_runs) == 0
def test_deploy_flow_must_be_implemented(cloud_api):
agent = Agent()
with pytest.raises(NotImplementedError):
agent.deploy_flow(None)
def test_get_flow_run_metadata(monkeypatch, cloud_api):
Client = MagicMock()
monkeypatch.setattr("prefect.agent.agent.Client", Client)
now = pendulum.now()
monkeypatch.setattr("prefect.agent.agent.pendulum.now", lambda *args: now)
agent = Agent()
agent._get_flow_run_metadata(["id1", "id2"])
Client().graphql.assert_called_with(
{
"query": {
with_args(
"flow_run",
{
"where": {
"id": {
"_in": ["id1", "id2"]
},
"_or": [
{
"state": {
"_eq": "Scheduled"
}
},
{
"state": {
"_eq": "Running"
},
"task_runs": {
"state_start_time": {
"_lte":
now.subtract(seconds=3).isoformat(
)
}
},
},
],
},
},
): {
"id": True,
"version": True,
"state": True,
"serialized_state": True,
"parameters": True,
"scheduled_start_time": True,
"run_config": True,
"name": True,
"flow": {
"version",
"id",
"environment",
"core_version",
"name",
"storage",
},
with_args(
"task_runs",
{
"where": {
"state_start_time": {
"_lte": now.subtract(seconds=3).isoformat(
)
}
}
},
): {
"serialized_state",
"version",
"id",
"task_id",
},
}
}
})
def test_agent_cloud_logs_responds_to_config_by_default(cloud_api, toggle):
with set_temporary_config(
{"cloud.agent.auth_token": "TEST_TOKEN", "cloud.send_flow_run_logs": toggle}
):
agent = Agent()
assert agent.log_to_cloud is toggle
def test_heartbeat_passes_base_agent(runner_token, cloud_api):
agent = Agent()
assert not agent.heartbeat()
def test_deploy_flows_passes_base_agent(runner_token):
agent = Agent()
assert not agent.deploy_flows([])
def test_on_flow_run_deploy_attempt_removes_id(monkeypatch, runner_token,
cloud_api):
agent = Agent()
agent.submitting_flow_runs.add("id")
agent.on_flow_run_deploy_attempt(None, "id")
assert len(agent.submitting_flow_runs) == 0
def test_agent_config_options(runner_token):
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = Agent()
assert agent.client.get_auth_token() == "TEST_TOKEN"
assert agent.logger
def test_agent_env_vars(runner_token, cloud_api):
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = Agent(env_vars=dict(AUTH_THING="foo"))
assert agent.env_vars == dict(AUTH_THING="foo")
def test_agent_fails_no_auth_token():
with pytest.raises(AuthorizationError):
agent = Agent()
agent.query_tenant_id()
def test_agent_max_polls(cloud_api):
with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
agent = Agent(max_polls=10)
assert agent.max_polls == 10
def test_agent_fails_no_auth_token(cloud_api):
with pytest.raises(AuthorizationError):
agent = Agent()
def test_agent_labels_from_config_var(cloud_api):
with set_temporary_config({"cloud.agent.labels": ["test", "2"]}):
agent = Agent()
assert agent.labels == ["test", "2"]
def test_agent_init(runner_token, cloud_api):
agent = Agent()
assert agent
def test_agent_init(cloud_api):
agent = Agent()
assert agent
def test_heartbeat_passes_base_agent(cloud_api):
agent = Agent()
assert not agent.heartbeat()