def test_docker_agent_is_named_volume_win32(monkeypatch, cloud_api, path, result): api = MagicMock() monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) monkeypatch.setattr("prefect.agent.docker.agent.platform", "win32") agent = DockerAgent() assert agent._is_named_volume_win32(path) == result
def test_docker_agent_parse_volume_spec_raises_on_invalid_spec( monkeypatch, cloud_api, candidate, exception_type): api = MagicMock() monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) agent = DockerAgent() with pytest.raises(exception_type): agent._parse_volume_spec([candidate])
def test_docker_agent_is_named_volume_unix(monkeypatch, runner_token, path, result): api = MagicMock() monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=api), ) monkeypatch.setattr("prefect.agent.docker.agent.platform", "osx") agent = DockerAgent() assert agent._is_named_volume_unix(path) == result
def test_docker_agent_shutdown_terminates_child_processes( monkeypatch, runner_token): 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.docker.APIClient", 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_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_init(monkeypatch, runner_token): api = MagicMock() monkeypatch.setattr("prefect.agent.docker.agent.docker.APIClient", api) agent = DockerAgent() assert agent assert agent.labels == [] assert agent.name == "agent"
def test_docker_agent_ping(monkeypatch, runner_token): api = MagicMock() api.ping.return_value = True monkeypatch.setattr("prefect.agent.docker.agent.docker.APIClient", MagicMock(return_value=api)) agent = DockerAgent() assert api.ping.called
def test_docker_agent_ping_exception(monkeypatch, runner_token): api = MagicMock() api.ping.return_value = True api.ping.side_effect = Exception() monkeypatch.setattr("prefect.agent.docker.agent.docker.APIClient", MagicMock(return_value=api)) with pytest.raises(Exception): agent = DockerAgent()
def test_docker_agent_heartbeat_exits_on_failure(monkeypatch, runner_token, caplog): api = MagicMock() api.ping.return_value = True monkeypatch.setattr("prefect.agent.docker.agent.docker.APIClient", 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_populate_env_vars_is_responsive_to_logging_config( monkeypatch, runner_token, flag): api = MagicMock() monkeypatch.setattr("prefect.agent.docker.agent.docker.APIClient", api) with set_temporary_config({ "cloud.agent.auth_token": "token", "cloud.api": "api", "logging.log_to_cloud": flag, }): agent = DockerAgent(labels=["42", "marvin"]) env_vars = agent.populate_env_vars( GraphQLResult({ "id": "id", "name": "name" })) assert env_vars["PREFECT__LOGGING__LOG_TO_CLOUD"] == str(flag).lower()
def test_populate_env_vars_uses_user_provided_env_vars(monkeypatch, runner_token): api = MagicMock() monkeypatch.setattr("prefect.agent.docker.agent.docker.APIClient", 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" })) assert env_vars["AUTH_THING"] == "foo"
def test_docker_agent_config_options_populated(monkeypatch, runner_token): api = MagicMock() monkeypatch.setattr("prefect.agent.docker.agent.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"
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_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_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_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_docker_agent_daemon_url_responds_to_system(monkeypatch, runner_token): api = MagicMock() monkeypatch.setattr("prefect.agent.docker.agent.docker.APIClient", api) monkeypatch.setattr("prefect.agent.docker.agent.platform", "win32") with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}): agent = DockerAgent() 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_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_deploy_flow_storage_raises(monkeypatch, runner_token): 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.docker.APIClient", MagicMock(return_value=api)) agent = DockerAgent() with pytest.raises(ValueError): agent.deploy_flow(flow_run=GraphQLResult( { "flow": GraphQLResult({"storage": Local().serialize()}), "id": "id", "name": "name", "version": "version", })) assert not api.pull.called
def test_docker_agent_no_pull(monkeypatch, runner_token): api = MagicMock() monkeypatch.setattr("prefect.agent.docker.agent.docker.APIClient", 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(monkeypatch, runner_token): api = MagicMock() monkeypatch.setattr("prefect.agent.docker.agent.docker.APIClient", api) monkeypatch.setattr("prefect.agent.docker.agent.platform", "osx") with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}): 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"] == "unix://var/run/docker.sock"
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_docker_agent_deploy_flow_reg_allow_list_not_allowed( monkeypatch, cloud_api): 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(reg_allow_list=["test1"]) with pytest.raises(ValueError) as error: agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "id": "foo", "storage": Docker(registry_url="test2", image_name="name", image_tag="tag").serialize(), "environment": LocalEnvironment().serialize(), "core_version": "0.13.0", }), "id": "id", "name": "name", })) expected_error = ("Trying to pull image from a Docker registry 'test2'" " which is not in the reg_allow_list") assert not api.pull.called assert not api.create_container.called assert not api.start.called assert str(error.value) == expected_error
def test_docker_agent_deploy_flow(core_version, command, 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": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "environment": LocalEnvironment().serialize(), "core_version": core_version, }), "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"] == command assert api.create_container.call_args[1]["host_config"][ "AutoRemove"] is True assert api.start.call_args[1]["container"] == "container_id"
def test_populate_env_vars(monkeypatch, runner_token): api = MagicMock() monkeypatch.setattr("prefect.agent.docker.agent.docker.APIClient", api) with set_temporary_config({ "cloud.agent.auth_token": "token", "cloud.api": "api", "logging.log_to_cloud": True, }): agent = DockerAgent() env_vars = agent.populate_env_vars( GraphQLResult({ "id": "id", "name": "name" })) expected_vars = { "PREFECT__CLOUD__API": "api", "PREFECT__CLOUD__AUTH_TOKEN": "token", "PREFECT__CLOUD__AGENT__LABELS": "[]", "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_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.assert_called_with(target=agent.stream_container_logs, kwargs={"container_id": "container_id"}) assert len(agent.processes) == 1 assert api.create_container.called assert api.start.called
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_docker_agent_no_pull(monkeypatch, cloud_api): api = MagicMock() monkeypatch.setattr( "prefect.agent.docker.agent.DockerAgent._get_docker_client", MagicMock(return_value=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, 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"
def test_docker_agent_heartbeat_logs_reconnect(monkeypatch, runner_token, caplog): api = MagicMock() api.ping.return_value = True monkeypatch.setattr("prefect.agent.docker.agent.docker.APIClient", MagicMock(return_value=api)) agent = DockerAgent() api.ping.return_value = False agent.heartbeat() agent.heartbeat() api.ping.return_value = True agent.heartbeat() assert api.ping.call_count == 4 assert "Reconnected to Docker daemon" in caplog.text