def test_local_agent_deploy_flows(monkeypatch): 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)) with set_temporary_config({"cloud.agent.auth_token": "token"}): agent = LocalAgent() agent.deploy_flows(flow_runs=[ GraphQLResult({ "flow": GraphQLResult({ "storage": Docker( registry_url="test", image_name="name", image_tag="tag", ).serialize() }), "id": "id", }) ]) assert api.pull.called assert api.create_container.called assert api.start.called assert (api.create_container.call_args[1]["command"] == "prefect execute cloud-flow") assert api.start.call_args[1]["container"] == "container_id"
def test_local_agent_uses_ip_if_dockerdesktop_hostname(monkeypatch): monkeypatch.setattr("socket.gethostname", MagicMock(return_value="docker-desktop")) monkeypatch.setattr("socket.gethostbyname", MagicMock(return_value="IP")) agent = LocalAgent() assert "IP" in agent.labels
def test_local_agent_deploy_no_existing_python_path(monkeypatch, cloud_api): 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": "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"]
def test_local_agent_init(monkeypatch, runner_token): api = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", api) agent = LocalAgent() assert agent assert agent.name == "agent"
def test_populate_env_vars(monkeypatch): 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( GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}}) ) expected = os.environ.copy() expected.update( { "PYTHONPATH": os.getcwd() + os.pathsep + expected.get("PYTHONPATH", ""), "PREFECT__CLOUD__API": "https://api.prefect.io", "PREFECT__CLOUD__AUTH_TOKEN": "TEST_TOKEN", "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__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
def test_local_agent_deploy_storage_raises_not_supported_storage( monkeypatch, cloud_api ): 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", "core_version": "0.13.0", }, }, ) ) assert not popen.called assert len(agent.processes) == 0
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_run_config_working_dir(monkeypatch, working_dir, tmpdir): popen = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.Popen", popen) if working_dir is not None: working_dir = str(tmpdir) agent = LocalAgent() 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 popen.called assert len(agent.processes) == 1 assert popen.call_args[1]["cwd"] == working_dir
def test_local_agent_deploy_unsupported_run_config(monkeypatch): popen = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.Popen", popen) agent = LocalAgent() with pytest.raises( TypeError, match="`run_config` of type `KubernetesRun`, only `LocalRun` is supported", ): agent.deploy_flow( flow_run=GraphQLResult( { "id": "id", "flow": { "storage": Local().serialize(), "id": "foo", "core_version": "0.13.0", }, "run_config": KubernetesRun().serialize(), }, ) ) assert not popen.called assert len(agent.processes) == 0
def test_local_agent_deploy_flows_no_registry_does_not_pull( 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": Docker(registry_url="", image_name="name", image_tag="tag").serialize() }), "id": "id", }) ]) assert not api.pull.called assert api.create_container.called assert api.start.called
def test_local_agent_deploy_storage_fails_none(monkeypatch): 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": "foo", "core_version": "0.13.0" }), "id": "id", "version": 1, })) assert not popen.called assert len(agent.processes) == 0 assert client.called
def test_populate_env_vars(runner_token): with set_temporary_config( { "cloud.api": "api", "logging.log_to_cloud": True, "cloud.agent.auth_token": "token", } ): agent = LocalAgent() env_vars = agent.populate_env_vars(GraphQLResult({"id": "id"})) expected_vars = { "PREFECT__CLOUD__API": "api", "PREFECT__CLOUD__AUTH_TOKEN": "token", "PREFECT__CLOUD__AGENT__LABELS": str( [ socket.gethostname(), "azure-flow-storage", "gcs-flow-storage", "s3-flow-storage", ] ), "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_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_deploy_processes_webhook_storage(monkeypatch, runner_token): popen = MagicMock() monkeypatch.setattr("prefect.agent.local.agent.Popen", popen) agent = LocalAgent() webhook = Webhook( build_request_kwargs={"url": "test-service/upload"}, build_request_http_method="POST", get_flow_request_kwargs={"url": "test-service/download"}, get_flow_request_http_method="GET", ) agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "storage": webhook.serialize(), "id": "foo", "core_version": "0.13.0", }), "id": "id", })) assert popen.called assert len(agent.processes) == 1
def test_local_agent_config_options(cloud_api): with set_temporary_config({ "cloud.agent.auth_token": "TEST_TOKEN", "logging.log_to_cloud": True }): agent = LocalAgent( name="test", labels=["test_label"], import_paths=["test_path"], hostname_label=False, ) 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) == { "azure-flow-storage", "s3-flow-storage", "gcs-flow-storage", "github-flow-storage", "webhook-flow-storage", "test_label", }
def test_populate_env_vars_sets_log_to_cloud(flag): agent = LocalAgent(no_cloud_logs=flag) assert agent.log_to_cloud is not flag env_vars = agent.populate_env_vars( GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}}) ) assert env_vars["PREFECT__LOGGING__LOG_TO_CLOUD"] == str(not flag).lower()
def test_populate_env_vars_from_run_config(tmpdir): agent = LocalAgent(env_vars={"KEY1": "VAL1", "KEY2": "VAL2"}) working_dir = str(tmpdir) run = LocalRun( env={ "KEY2": "OVERRIDE", "PREFECT__LOGGING__LEVEL": "TEST" }, working_dir=working_dir, ) env_vars = agent.populate_env_vars( GraphQLResult({ "id": "id", "name": "name", "flow": { "id": "foo" }, "run_config": run.serialize(), }), run, ) assert env_vars["KEY1"] == "VAL1" assert env_vars["KEY2"] == "OVERRIDE" assert env_vars["PREFECT__LOGGING__LEVEL"] == "TEST" assert working_dir in env_vars["PYTHONPATH"]
def test_populate_env_vars_includes_agent_labels(): agent = LocalAgent(labels=["42", "marvin"]) env_vars = agent.populate_env_vars( GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}}) ) expected = str(["42", "marvin"] + DEFAULT_AGENT_LABELS) assert env_vars["PREFECT__CLOUD__AGENT__LABELS"] == expected
def test_populate_env_vars_removes_none_values(): 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_responds_to_logging_config(runner_token, flag): with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}): agent = LocalAgent(no_cloud_logs=flag) assert agent.log_to_cloud is not flag env_vars = agent.populate_env_vars( GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}}) ) assert env_vars["PREFECT__LOGGING__LOG_TO_CLOUD"] == str(not flag).lower()
def test_populate_env_vars_no_existing_python_path(monkeypatch): monkeypatch.delenv("PYTHONPATH", raising=False) agent = LocalAgent(import_paths=["paths"]) env_vars = agent.populate_env_vars( GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}}) ) assert "paths" in env_vars["PYTHONPATH"]
def test_populate_env_vars_from_agent_config(): agent = LocalAgent(env_vars=dict(AUTH_THING="foo")) env_vars = agent.populate_env_vars( GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}}) ) assert env_vars["AUTH_THING"] == "foo"
def test_environment_has_agent_token_from_config(): """Check that the API token is passed through from the config via environ""" with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}): agent = LocalAgent() env = agent.populate_env_vars(TEST_FLOW_RUN_DATA) assert env["PREFECT__CLOUD__AUTH_TOKEN"] == "TEST_TOKEN"
def test_local_agent_ping(monkeypatch, runner_token): api = MagicMock() api.ping.return_value = True monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", MagicMock(return_value=api)) agent = LocalAgent() assert api.ping.called
def test_local_agent_heartbeat( monkeypatch, cloud_api, returncode, show_flow_logs, logs ): popen = MockPopen() # expect a process to be called with the following command (with specified behavior) popen.set_command( "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 test_local_agent_ping(monkeypatch): api = MagicMock() api.ping.return_value = True monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient", MagicMock(return_value=api)) with set_temporary_config({"cloud.agent.auth_token": "token"}): agent = LocalAgent() assert api.ping.called
def test_local_agent_config_no_storage_labels(): agent = LocalAgent( labels=["test_label"], storage_labels=False, ) assert set(agent.labels) == { socket.gethostname(), "test_label", }
def test_populate_env_vars_keep_existing_python_path(monkeypatch): monkeypatch.setenv("PYTHONPATH", "cool:python:path") agent = LocalAgent(import_paths=["paths"]) env_vars = agent.populate_env_vars(TEST_FLOW_RUN_DATA) python_path = env_vars["PYTHONPATH"] assert "cool:python:path" in python_path assert "paths" in python_path
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"