def test_local_agent_init(runner_token): agent = LocalAgent() assert agent assert set(agent.labels) == { socket.gethostname(), "azure-flow-storage", "s3-flow-storage", "gcs-flow-storage", } assert agent.name == "agent"
def test_local_agent_deploy_raises_unsupported_storage(monkeypatch): popen = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.Popen", popen) agent = LocalAgent() with pytest.raises(TypeError, match="Unsupported Storage type: Docker"): agent.deploy_flow(flow_run=GraphQLResult( { "id": "id", "flow": { "storage": Docker().serialize(), "id": "foo", "core_version": "0.13.0", }, }, )) assert not popen.called assert len(agent.processes) == 0
def test_populate_env_vars_uses_user_provided_env_vars_removes_nones( runner_token): with set_temporary_config({ "cloud.api": "api", "logging.log_to_cloud": True, "cloud.agent.auth_token": "token", }): agent = LocalAgent(env_vars=dict(MISSING_VAR=None)) env_vars = agent.populate_env_vars( GraphQLResult({ "id": "id", "name": "name", "flow": { "id": "foo" } })) assert "MISSING_VAR" not in env_vars
def test_local_agent_start_max_polls_zero(monkeypatch, runner_token): on_shutdown = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.LocalAgent.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.local.agent.LocalAgent.heartbeat", heartbeat) agent = LocalAgent(max_polls=0) agent.start() assert on_shutdown.call_count == 1 assert agent_process.call_count == 0 assert heartbeat.call_count == 0
def test_local_agent_config_options_hostname(runner_token): with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}): agent = LocalAgent(labels=["test_label"]) assert set(agent.labels) == { "test_label", socket.gethostname(), "azure-flow-storage", "s3-flow-storage", "gcs-flow-storage", }
def test_local_agent_deploy_no_existing_python_path(monkeypatch, runner_token): monkeypatch.delenv("PYTHONPATH", raising=False) popen = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.Popen", popen) agent = LocalAgent(import_paths=["paths"]) agent.deploy_flow( flow_run=GraphQLResult( { "flow": GraphQLResult({"storage": Local(directory="test").serialize()}), "id": "id", } ) ) assert popen.called assert len(agent.processes) == 1 assert "paths" in popen.call_args[1]["env"]["PYTHONPATH"]
def test_local_agent_ping_exception(monkeypatch, runner_token): api = MagicMock() api.ping.return_value = True api.ping.side_effect = Exception() monkeypatch.setattr( "prefect.agent.local.agent.docker.APIClient", MagicMock(return_value=api) ) with pytest.raises(Exception): agent = LocalAgent()
def test_local_agent_config_options_populated(monkeypatch, runner_token): api = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api) with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}): agent = LocalAgent(base_url="url", no_pull=True) assert agent.client.get_auth_token() == "TEST_TOKEN" assert agent.logger assert agent.no_pull assert api.call_args[1]["base_url"] == "url"
def test_local_agent_deploy_flows_storage_continues(monkeypatch, runner_token): api = MagicMock() api.ping.return_value = True api.create_container.return_value = {"Id": "container_id"} monkeypatch.setattr( "prefect.agent.local.agent.docker.APIClient", MagicMock(return_value=api) ) agent = LocalAgent() agent.deploy_flows( flow_runs=[ GraphQLResult( {"flow": GraphQLResult({"storage": Local().serialize()}), "id": "id"} ) ] ) assert not api.pull.called
def test_local_agent_heartbeat(monkeypatch, returncode, show_flow_logs, logs): popen = MockPopen() # expect a process to be called with the following command (with specified behavior) popen.set_command( [sys.executable, "-m", "prefect", "execute", "flow-run"], stdout=b"awesome output!", stderr=b"blerg, eRroR!", returncode=returncode, poll_count=2, ) monkeypatch.setattr("prefect.agent.local.agent.Popen", popen) agent = LocalAgent(import_paths=["paths"], show_flow_logs=show_flow_logs) agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "storage": Local(directory="test").serialize(), "id": "foo", "core_version": "0.13.0", }), "id": "id", })) process = list(agent.processes)[0] process_call = process.root_call with LogCapture() as logcap: agent.heartbeat() agent.heartbeat() agent.heartbeat() # ensure the expected logs exist (or the absense of logs) if logs: logcap.check(*logs) else: logcap.check() # ensure the process was opened and was polled compare( popen.all_calls, expected=[ process_call, process_call.poll(), process_call.poll(), process_call.poll(), ], ) # the heartbeat should stop tracking upon exit compare(process.returncode, returncode) assert len(agent.processes) == 0
def install(label, env, import_paths, **kwargs): """Generate a supervisord.conf file for a Local agent""" from prefect.agent.local import LocalAgent conf = LocalAgent.generate_supervisor_conf( labels=sorted(set(label)), env_vars=dict(e.split("=", 1) for e in env), import_paths=list(import_paths), **kwargs, ) click.echo(conf)
def test_local_agent_config_options(monkeypatch, runner_token): api = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api) monkeypatch.setattr("prefect.agent.local.agent.platform", "osx") with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}): agent = LocalAgent() assert agent.client.get_auth_token() == "TEST_TOKEN" assert agent.logger assert not agent.no_pull assert api.call_args[1]["base_url"] == "unix://var/run/docker.sock"
def test_local_agent_daemon_url_responds_to_system(monkeypatch, runner_token): api = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api) monkeypatch.setattr("prefect.agent.local.agent.platform", "win32") with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}): agent = LocalAgent() assert agent.client.get_auth_token() == "TEST_TOKEN" assert agent.logger assert not agent.no_pull assert api.call_args[1]["base_url"] == "npipe:////./pipe/docker_engine"
def test_local_agent_deploy_null_or_univeral_run_config( monkeypatch, run_config): popen = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.Popen", popen) agent = LocalAgent() agent.deploy_flow(flow_run=GraphQLResult( { "id": "id", "flow": { "storage": Local().serialize(), "id": "foo", "core_version": "0.13.0", }, "run_config": run_config.serialize() if run_config else None, }, )) assert popen.called assert len(agent.processes) == 1
def test_local_agent_deploy_storage_raises_not_supported_storage( monkeypatch, runner_token): popen = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.Popen", popen) agent = LocalAgent() with pytest.raises(ValueError): agent.deploy_flow(flow_run=GraphQLResult( { "id": "id", "flow": { "storage": Docker().serialize(), "id": "foo" } }, )) assert not popen.called assert len(agent.processes) == 0
def test_local_agent_deduplicates_labels(runner_token): agent = LocalAgent(labels=["azure-flow-storage"]) assert agent assert set(agent.labels) == { socket.gethostname(), "azure-flow-storage", "s3-flow-storage", "gcs-flow-storage", "github-flow-storage", } assert len(agent.labels) == len(set(agent.labels))
def test_populate_env_vars(monkeypatch, backend, config_with_api_key): agent = LocalAgent() # The python path may be a single item and we want to ensure the correct separator # is added so we will ensure PYTHONPATH has an item in it to start if not os.environ.get("PYTHONPATH", ""): monkeypatch.setenv("PYTHONPATH", "foobar") env_vars = agent.populate_env_vars(TEST_FLOW_RUN_DATA) expected = os.environ.copy() expected.update({ "PYTHONPATH": os.getcwd() + os.pathsep + expected.get("PYTHONPATH", ""), "PREFECT__BACKEND": backend, "PREFECT__CLOUD__API": prefect.config.cloud.api, "PREFECT__CLOUD__API_KEY": config_with_api_key.cloud.api_key, "PREFECT__CLOUD__TENANT_ID": config_with_api_key.cloud.tenant_id, "PREFECT__CLOUD__AGENT__LABELS": str(DEFAULT_AGENT_LABELS), "PREFECT__CONTEXT__FLOW_RUN_ID": "id", "PREFECT__CONTEXT__FLOW_ID": "foo", "PREFECT__CLOUD__USE_LOCAL_SECRETS": "false", "PREFECT__CLOUD__SEND_FLOW_RUN_LOGS": "true", "PREFECT__LOGGING__LEVEL": "INFO", "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudFlowRunner", "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudTaskRunner", }) assert env_vars == expected
def test_populate_env_vars(monkeypatch, runner_token): api = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api) with set_temporary_config({"cloud.agent.auth_token": "token", "cloud.api": "api"}): agent = LocalAgent() env_vars = agent.populate_env_vars(GraphQLResult({"id": "id"})) expected_vars = { "PREFECT__CLOUD__API": "api", "PREFECT__CLOUD__AUTH_TOKEN": "token", "PREFECT__CONTEXT__FLOW_RUN_ID": "id", "PREFECT__CLOUD__USE_LOCAL_SECRETS": "false", "PREFECT__LOGGING__LOG_TO_CLOUD": "true", "PREFECT__LOGGING__LEVEL": "DEBUG", "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudFlowRunner", "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudTaskRunner", } assert env_vars == expected_vars
def test_local_agent_init(cloud_api): agent = LocalAgent() assert agent assert set(agent.labels) == { socket.gethostname(), "azure-flow-storage", "s3-flow-storage", "gcs-flow-storage", "github-flow-storage", "webhook-flow-storage", } assert agent.name == "agent"
def test_local_agent_no_pull(monkeypatch, runner_token): api = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api) agent = LocalAgent() assert not agent.no_pull agent = LocalAgent(no_pull=True) assert agent.no_pull with context(no_pull=True): agent = LocalAgent() assert agent.no_pull with context(no_pull=False): agent = LocalAgent(no_pull=True) assert agent.no_pull with context(no_pull=False): agent = LocalAgent(no_pull=False) assert not agent.no_pull
def test_populate_env_vars_includes_agent_labels(cloud_api): with set_temporary_config( { "cloud.api": "api", "logging.log_to_cloud": True, "cloud.agent.auth_token": "token", } ): agent = LocalAgent(labels=["42", "marvin"]) env_vars = agent.populate_env_vars( GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}}) ) expected_vars = { "PREFECT__CLOUD__API": "api", "PREFECT__CLOUD__AUTH_TOKEN": "token", "PREFECT__CLOUD__AGENT__LABELS": str( [ "42", "marvin", socket.gethostname(), "azure-flow-storage", "gcs-flow-storage", "s3-flow-storage", "github-flow-storage", "webhook-flow-storage", "gitlab-flow-storage", ] ), "PREFECT__CONTEXT__FLOW_RUN_ID": "id", "PREFECT__CONTEXT__FLOW_ID": "foo", "PREFECT__CLOUD__USE_LOCAL_SECRETS": "false", "PREFECT__LOGGING__LOG_TO_CLOUD": "true", "PREFECT__LOGGING__LEVEL": "INFO", "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudFlowRunner", "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudTaskRunner", } assert env_vars == expected_vars
def test_local_agent_config_options(config_with_api_key): agent = LocalAgent( name="test", labels=["test_label"], import_paths=["test_path"], ) assert agent.name == "test" assert agent.client.api_key == config_with_api_key.cloud.api_key assert agent.logger assert agent.log_to_cloud is True assert agent.processes == set() assert agent.import_paths == ["test_path"] assert set(agent.labels) == {"test_label", *DEFAULT_AGENT_LABELS}
def test_local_agent_config_options(): agent = LocalAgent( name="test", labels=["test_label"], import_paths=["test_path"], ) assert agent.name == "test" assert agent.client.get_auth_token() == "TEST_TOKEN" assert agent.logger assert agent.log_to_cloud is True assert agent.processes == set() assert agent.import_paths == ["test_path"] assert set(agent.labels) == {"test_label", *DEFAULT_AGENT_LABELS}
def test_local_agent_deploy_processes_gitlab_storage(monkeypatch, cloud_api): popen = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.Popen", popen) agent = LocalAgent() gitlab = GitLab("test/repo", path="path/to/flow.py") agent.deploy_flow( flow_run=GraphQLResult( { "flow": GraphQLResult( { "storage": gitlab.serialize(), "id": "foo", } ), "id": "id", } ) ) assert popen.called assert len(agent.processes) == 1
def test_local_agent_no_pull(monkeypatch): api = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api) with set_temporary_config({"cloud.agent.auth_token": "token"}): agent = LocalAgent() assert not agent.no_pull agent = LocalAgent(no_pull=True) assert agent.no_pull with context(no_pull=True): agent = LocalAgent() assert agent.no_pull with context(no_pull=False): agent = LocalAgent(no_pull=True) assert agent.no_pull with context(no_pull=False): agent = LocalAgent(no_pull=False) assert not agent.no_pull
def test_local_agent_deploy_processes_valid_storage(storage, monkeypatch): popen = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.Popen", popen) agent = LocalAgent() agent.deploy_flow( flow_run=GraphQLResult( { "flow": GraphQLResult( { "storage": storage.serialize(), "id": "foo", "core_version": "0.13.0", } ), "id": "id", } ) ) assert popen.called assert len(agent.processes) == 1
def test_local_agent_config_no_storage_labels(cloud_api): with set_temporary_config({ "cloud.agent.auth_token": "TEST_TOKEN", "logging.log_to_cloud": True }): agent = LocalAgent( labels=["test_label"], storage_labels=False, ) assert set(agent.labels) == { socket.gethostname(), "test_label", }
def test_local_agent_deploy_storage_fails_none(monkeypatch, runner_token): client = MagicMock() set_state = MagicMock() client.return_value.set_flow_run_state = set_state monkeypatch.setattr("prefect.agent.agent.Client", client) popen = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.Popen", popen) agent = LocalAgent() with pytest.raises(ValidationError): agent.deploy_flow( flow_run=GraphQLResult( {"flow": GraphQLResult({"storage": None}), "id": "id", "version": 1} ) ) assert not popen.called assert len(agent.processes) == 0 assert client.called
def test_local_agent_deploy_run_config_missing_working_dir( monkeypatch, tmpdir): popen = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.Popen", popen) working_dir = str(tmpdir.join("missing")) agent = LocalAgent() with pytest.raises(ValueError, match="nonexistent `working_dir`"): agent.deploy_flow(flow_run=GraphQLResult( { "id": "id", "flow": { "storage": Local().serialize(), "id": "foo", "core_version": "0.13.0", }, "run_config": LocalRun(working_dir=working_dir).serialize(), }, )) assert not popen.called assert not agent.processes
def test_local_agent_deploy_import_paths(monkeypatch, cloud_api): popen = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.Popen", popen) agent = LocalAgent(import_paths=["paths"]) agent.deploy_flow( flow_run=GraphQLResult( { "flow": GraphQLResult( { "storage": Local(directory="test").serialize(), "id": "foo", "core_version": "0.13.0", } ), "id": "id", } ) ) assert popen.called assert len(agent.processes) == 1 assert "paths" in popen.call_args[1]["env"]["PYTHONPATH"]