def test_docker_agent_parse_volume_spec_win(api, candidate, named_volumes, container_mount_paths, host_spec): agent = DockerAgent() ( actual_named_volumes, actual_container_mount_paths, actual_host_spec, ) = agent._parse_volume_spec_win32(candidate) assert actual_named_volumes == named_volumes assert actual_container_mount_paths == container_mount_paths assert actual_host_spec == host_spec
def test_environment_has_agent_token_from_config(api, config_with_token): agent = DockerAgent() env_vars = agent.populate_env_vars( GraphQLResult({ "id": "id", "name": "name", "flow": { "id": "foo" } }), "test-image") assert env_vars["PREFECT__CLOUD__AUTH_TOKEN"] == "TEST_TOKEN"
def test_populate_env_vars_from_agent_config(api): agent = DockerAgent(env_vars=dict(AUTH_THING="foo")) env_vars = agent.populate_env_vars( GraphQLResult({ "id": "id", "name": "name", "flow": { "id": "foo" } }), "test-image") assert env_vars["AUTH_THING"] == "foo"
def test_populate_env_vars(api, backend): agent = DockerAgent() env_vars = agent.populate_env_vars( GraphQLResult({ "id": "id", "name": "name", "flow": { "id": "foo" } }), "test-image") if backend == "server": cloud_api = "http://host.docker.internal:4200" else: cloud_api = prefect.config.cloud.api expected_vars = { "PREFECT__BACKEND": backend, "PREFECT__CLOUD__API": cloud_api, "PREFECT__CLOUD__AUTH_TOKEN": "", "PREFECT__CLOUD__API_KEY": "", "PREFECT__CLOUD__TENANT_ID": "", "PREFECT__CLOUD__AGENT__LABELS": "[]", "PREFECT__CONTEXT__FLOW_RUN_ID": "id", "PREFECT__CONTEXT__FLOW_ID": "foo", "PREFECT__CONTEXT__IMAGE": "test-image", "PREFECT__CLOUD__USE_LOCAL_SECRETS": "false", "PREFECT__CLOUD__SEND_FLOW_RUN_LOGS": "true", "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_api_url_can_be_overridden_at_agent_level(api): agent = DockerAgent(env_vars={"PREFECT__CLOUD__API": "FOO"}) env_vars = agent.populate_env_vars( GraphQLResult({ "id": "id", "name": "name", "flow": { "id": "foo" }, }), "test-image", ) assert env_vars["PREFECT__CLOUD__API"] == "FOO"
def test_docker_agent_deploy_flow_sets_container_name_with_index( api, collision_count): """ Asserts that the container name is set to the flow run name and that collisions with existing containers with the same name is handled by adding an index """ if collision_count: # Add the basic name first existing_names = ["flow-run-name"] for i in range(1, collision_count): existing_names.append(f"flow-run-name-{i}") else: existing_names = [] def fail_if_name_exists(*args, **kwargs): if kwargs.get("name") in existing_names: raise docker.errors.APIError( "Conflict. The container name 'foobar' is already in use") return {} api.create_container = MagicMock(side_effect=fail_if_name_exists) agent = DockerAgent() agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "id": "foo", "name": "flow-name", "storage": Local().serialize(), "environment": LocalEnvironment(metadata={ "image": "repo/name:tag" }).serialize(), "core_version": "0.13.0", }), "id": "id", "name": "flow-run-name", })) expected_name = ("flow-run-name" if not collision_count else f"flow-run-name-{collision_count}") assert api.create_container.call_args[1]["name"] == expected_name
def test_populate_env_vars_sets_log_to_cloud(flag, api, config_with_token): agent = DockerAgent(labels=["42", "marvin"], no_cloud_logs=flag) env_vars = agent.populate_env_vars( GraphQLResult({ "id": "id", "name": "name", "flow": { "id": "foo" } }), "test-image") assert env_vars["PREFECT__CLOUD__SEND_FLOW_RUN_LOGS"] == str( not flag).lower() # Backwards compatibility variable for containers on Prefect <0.15.0 assert env_vars["PREFECT__LOGGING__LOG_TO_CLOUD"] == str(not flag).lower()
def test_populate_env_vars_is_responsive_to_logging_config( monkeypatch, cloud_api, flag ): api = MagicMock() monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) with set_temporary_config({"cloud.agent.auth_token": "token", "cloud.api": "api"}): agent = DockerAgent(labels=["42", "marvin"], no_cloud_logs=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_includes_agent_labels(monkeypatch, cloud_api): api = MagicMock() monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) with set_temporary_config({ "cloud.agent.auth_token": "token", "cloud.api": "api", "logging.log_to_cloud": True, }): agent = DockerAgent(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__AGENT__LABELS": "['42', 'marvin']", "PREFECT__CLOUD__AUTH_TOKEN": "token", "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_docker_agent_deploy_flow_show_flow_logs(monkeypatch, cloud_api): process = MagicMock() monkeypatch.setattr("multiprocessing.Process", process) api = MagicMock() api.ping.return_value = True api.create_container.return_value = {"Id": "container_id"} monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) agent = DockerAgent(show_flow_logs=True) agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "id": "foo", "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "environment": LocalEnvironment().serialize(), "core_version": "0.13.0", }), "id": "id", "name": "name", })) process_kwargs = dict( target=_stream_container_logs, kwargs={ "base_url": agent.base_url, "container_id": "container_id" }, ) process.assert_called_with(**process_kwargs) # Check all arguments to `multiprocessing.Process` are pickleable assert pickle.loads(pickle.dumps(process_kwargs)) == process_kwargs assert len(agent.processes) == 1 assert api.create_container.called assert api.start.called
def test_docker_agent_shutdown_terminates_child_processes(monkeypatch, cloud_api): monkeypatch.setattr("prefect.agent.agent.Client", MagicMock()) api = MagicMock() api.ping.return_value = True api.create_container.return_value = {"Id": "container_id"} monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) proc = MagicMock(is_alive=MagicMock(return_value=True)) agent = DockerAgent(show_flow_logs=True) agent.processes = [proc] agent.on_shutdown() assert proc.is_alive.called assert proc.terminate.called
def test_api_url_can_be_overridden_with_run_config(api): agent = DockerAgent(env_vars={"PREFECT__CLOUD__API": "FOO"}) run = DockerRun(env={"PREFECT__CLOUD__API": "BAR"}, ) env_vars = agent.populate_env_vars( GraphQLResult({ "id": "id", "name": "name", "flow": { "id": "foo" }, "run_config": run.serialize(), }), "test-image", run_config=run, ) assert env_vars["PREFECT__CLOUD__API"] == "BAR"
def test_docker_agent_no_pull(api): agent = DockerAgent() assert not agent.no_pull agent = DockerAgent(no_pull=True) assert agent.no_pull with context(no_pull=True): agent = DockerAgent() assert agent.no_pull with context(no_pull=False): agent = DockerAgent(no_pull=True) assert agent.no_pull with context(no_pull=False): agent = DockerAgent(no_pull=False) assert not agent.no_pull
def test_docker_agent_config_options_populated(monkeypatch): api = MagicMock() monkeypatch.setattr("docker.APIClient", api) agent = DockerAgent(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_docker_agent_ping(monkeypatch, cloud_api): api = MagicMock() api.ping.return_value = True monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) agent = DockerAgent() assert api.ping.called
def test_docker_agent_heartbeat_exits_on_failure(monkeypatch, cloud_api, caplog): api = MagicMock() api.ping.return_value = True monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) agent = DockerAgent() api.ping.return_value = False agent.heartbeat() agent.heartbeat() agent.heartbeat() agent.heartbeat() agent.heartbeat() with pytest.raises(SystemExit): agent.heartbeat() assert "Cannot reconnect to Docker daemon. Agent is shutting down." in caplog.text assert api.ping.call_count == 7
def test_docker_agent_deploy_flow_uses_environment_metadata( monkeypatch, cloud_api): api = MagicMock() api.ping.return_value = True api.create_container.return_value = {"Id": "container_id"} api.create_host_config.return_value = {"AutoRemove": True} monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) agent = DockerAgent() agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "id": "foo", "storage": Local().serialize(), "environment": LocalEnvironment(metadata={ "image": "repo/name:tag" }).serialize(), "core_version": "0.13.0", }), "id": "id", "name": "name", })) assert api.pull.called assert api.create_container.called assert api.start.called assert api.create_host_config.call_args[1]["auto_remove"] is True assert api.create_container.call_args[1][ "command"] == "prefect execute flow-run" assert api.create_container.call_args[1]["host_config"][ "AutoRemove"] is True assert api.start.call_args[1]["container"] == "container_id"
def test_environment_has_api_key_from_config(api, config_with_api_key): agent = DockerAgent() env_vars = agent.populate_env_vars( GraphQLResult({ "id": "id", "name": "name", "flow": { "id": "foo" } }), "test-image", ) assert env_vars[ "PREFECT__CLOUD__API_KEY"] == config_with_api_key.cloud.api_key assert env_vars[ "PREFECT__CLOUD__AUTH_TOKEN"] == config_with_api_key.cloud.api_key assert env_vars[ "PREFECT__CLOUD__TENANT_ID"] == config_with_api_key.cloud.tenant_id
def test_docker_agent_ping_exception(monkeypatch, cloud_api): api = MagicMock() api.ping.return_value = True api.ping.side_effect = Exception() monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) with pytest.raises(Exception): agent = DockerAgent()
def test_docker_agent_init(monkeypatch, cloud_api): api = MagicMock() monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) agent = DockerAgent() assert agent assert agent.labels == [] assert agent.name == "agent"
def test_docker_agent_config_options(platform, url, monkeypatch): api = MagicMock() monkeypatch.setattr("docker.APIClient", api) monkeypatch.setattr("prefect.agent.docker.agent.platform", platform) agent = DockerAgent(name="test") assert agent.name == "test" assert agent.client.get_auth_token() == "TEST_TOKEN" assert agent.logger assert not agent.no_pull assert api.call_args[1]["base_url"] == url
def test_docker_agent_parse_volume_spec_win( monkeypatch, cloud_api, candidate, named_volumes, container_mount_paths, host_spec ): api = MagicMock() monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) agent = DockerAgent() ( actual_named_volumes, actual_container_mount_paths, actual_host_spec, ) = agent._parse_volume_spec_win32(candidate) assert actual_named_volumes == named_volumes assert actual_container_mount_paths == container_mount_paths assert actual_host_spec == host_spec
def test_populate_env_vars_from_run_config(api): agent = DockerAgent(env_vars={"KEY1": "VAL1", "KEY2": "VAL2"}) run = DockerRun( env={"KEY2": "OVERRIDE", "PREFECT__LOGGING__LEVEL": "TEST"}, ) 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"
def test_docker_agent_deploy_flow_run_config(api, run_kind, has_docker_storage): if has_docker_storage: storage = Docker( registry_url="testing", image_name="on-storage", image_tag="tag" ) image = "testing/on-storage:tag" else: storage = Local() image = "on-run-config" if run_kind == "docker" else "prefecthq/prefect:0.13.11" if run_kind == "docker": env = {"TESTING": "VALUE"} run = DockerRun(image=image, env=env) else: env = {} run = None if run_kind == "missing" else UniversalRun() agent = DockerAgent() agent.deploy_flow( flow_run=GraphQLResult( { "flow": GraphQLResult( { "id": "foo", "name": "flow-name", "storage": storage.serialize(), "core_version": "0.13.11", } ), "run_config": run.serialize() if run else None, "id": "id", "name": "name", } ) ) assert api.create_container.called assert api.create_container.call_args[0][0] == image res_env = api.create_container.call_args[1]["environment"] for k, v in env.items(): assert res_env[k] == v
def test_docker_agent_config_options(platform, url, monkeypatch, config_with_api_key): api = MagicMock() monkeypatch.setattr("docker.APIClient", api) monkeypatch.setattr("prefect.agent.docker.agent.platform", platform) agent = DockerAgent(name="test") assert agent.name == "test" assert agent.client.api_key == config_with_api_key.cloud.api_key assert agent.logger assert not agent.no_pull assert api.call_args[1]["base_url"] == url
def test_populate_env_vars(api): agent = DockerAgent() env_vars = agent.populate_env_vars( GraphQLResult({"id": "id", "name": "name", "flow": {"id": "foo"}}) ) expected_vars = { "PREFECT__CLOUD__API": "https://api.prefect.io", "PREFECT__CLOUD__AUTH_TOKEN": "TEST_TOKEN", "PREFECT__CLOUD__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_vars
def test_populate_env_vars_uses_user_provided_env_vars(monkeypatch, cloud_api): api = MagicMock() monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) with set_temporary_config( { "cloud.agent.auth_token": "token", "cloud.api": "api", "logging.log_to_cloud": True, } ): agent = DockerAgent(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_docker_agent_init_volume_empty_options(monkeypatch, cloud_api): api = MagicMock() monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) agent = DockerAgent() assert agent assert agent.named_volumes == [] assert agent.container_mount_paths == [] assert agent.host_spec == {}
def test_prefect_logging_level_override_logic(config, agent_env_vars, run_config_env_vars, expected_logging_level, api): with set_temporary_config(config): agent = DockerAgent(env_vars=agent_env_vars) run = DockerRun(env=run_config_env_vars) env_vars = agent.populate_env_vars( GraphQLResult({ "id": "id", "name": "name", "flow": { "id": "foo" }, "run_config": run.serialize(), }), "test-image", run_config=run, ) assert env_vars["PREFECT__LOGGING__LEVEL"] == expected_logging_level
def test_docker_agent_config_options_populated(monkeypatch, cloud_api): import docker # DockerAgent imports docker within the constructor api = MagicMock() monkeypatch.setattr("docker.APIClient", api) with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}): agent = DockerAgent(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"