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