def test_docker_agent_deploy_flow_no_registry_does_not_pull(api): agent = DockerAgent() agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "id": "foo", "name": "flow-name", "storage": Docker(registry_url="", image_name="name", image_tag="tag").serialize(), "environment": LocalEnvironment().serialize(), "core_version": "0.13.0", }), "id": "id", "name": "name", })) assert not api.pull.called assert api.create_container.called assert api.start.called
def test_initialized_docker_storage(no_docker_host_var): storage = Docker( registry_url="test1", base_image="test3", python_dependencies=["test"], image_name="test4", image_tag="test5", env_vars={"test": "1"}, base_url="test_url", tls_config={"tls": "here"}, prefect_version="my-branch", local_image=True, build_kwargs={"nocache": True}, ) assert storage.registry_url == "test1" assert storage.base_image == "test3" assert storage.image_name == "test4" assert storage.image_tag == "test5" assert storage.python_dependencies == ["test", "wheel"] assert storage.env_vars == { "test": "1", "PREFECT__USER_CONFIG_PATH": "/opt/prefect/config.toml", } assert storage.base_url == "test_url" assert storage.tls_config == {"tls": "here"} assert storage.build_kwargs == {"nocache": True} assert storage.prefect_version == "my-branch" assert storage.local_image
def test_pull_image_raises_if_error_encountered(monkeypatch): storage = Docker(base_image="python:3.6") client = MagicMock() client.pull.return_value = [ { "progress": "test" }, { "error": "you know nothing jon snow" }, ] monkeypatch.setattr("docker.APIClient", MagicMock(return_value=client)) with pytest.raises(InterruptedError, match="you know nothing jon snow"): storage.pull_image()
def test_docker_agent_deploy_with_interface_check_linux( api, monkeypatch, linux_platform): get_ip = MagicMock() monkeypatch.setattr("prefect.agent.docker.agent.get_docker_ip", get_ip) agent = DockerAgent() agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "id": "foo", "name": "flow-name", "storage": Docker(registry_url="", image_name="name", image_tag="tag").serialize(), "environment": LocalEnvironment().serialize(), "core_version": "0.13.0", }), "id": "id", "name": "name", })) assert get_ip.called
def test_docker_agent_network(api): api.create_networking_config.return_value = {"test-network": "config"} with pytest.warns(UserWarning): agent = DockerAgent(network="test-network") agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "id": "foo", "name": "flow-name", "storage": Docker(registry_url="test", image_name="name", image_tag="tag").serialize(), "environment": LocalEnvironment().serialize(), "core_version": "0.13.0", }), "id": "id", "name": "name", })) assert agent.network == "test-network" assert agent.networks is None args, kwargs = api.create_container.call_args assert kwargs["networking_config"] == {"test-network": "config"}
def test_docker_agent_deploy_flow_reg_allow_list_not_allowed(api): agent = DockerAgent(reg_allow_list=["test1"]) with pytest.raises(ValueError) as error: agent.deploy_flow(flow_run=GraphQLResult({ "flow": GraphQLResult({ "id": "foo", "name": "flow-name", "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_create_dockerfile_from_dockerfile_uses_tempdir_path(): myfile = "FROM my-own-image:latest\n\nRUN echo 'hi'" with tempfile.TemporaryDirectory() as tempdir_outside: with open(os.path.join(tempdir_outside, "test"), "w+") as t: t.write("asdf") with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "myfile"), "w") as tmp: tmp.write(myfile) storage = Docker( dockerfile=os.path.join(directory, "myfile"), files={os.path.join(tempdir_outside, "test"): "./test2"}, ) storage.add_flow(Flow("foo")) dpath = storage.create_dockerfile_object(directory=directory) with open(dpath, "r") as dockerfile: output = dockerfile.read() assert ( "COPY {} /opt/prefect/flows/foo.prefect".format( os.path.join(directory, "foo.flow").replace("\\", "/") ) in output ), output assert ( "COPY {} ./test2".format( os.path.join(directory, "test").replace("\\", "/") ) in output ), output assert ( "COPY {} /opt/prefect/healthcheck.py".format( os.path.join(directory, "healthcheck.py").replace("\\", "/") ) in output ) assert output.startswith("\n" + myfile) # test proper indentation assert all( line == line.lstrip() for line in output.split("\n") if line not in ["\n", " "] )
def test_setup_definition_exists(monkeypatch): existing_task_definition = { "containerDefinitions": [{ "environment": [ { "name": "PREFECT__CLOUD__GRAPHQL", "value": config.cloud.graphql }, { "name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false" }, { "name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS", "value": "prefect.engine.cloud.CloudFlowRunner", }, { "name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS", "value": "prefect.engine.cloud.CloudTaskRunner", }, { "name": "PREFECT__CLOUD__SEND_FLOW_RUN_LOGS", "value": "true" }, { "name": "PREFECT__LOGGING__EXTRA_LOGGERS", "value": str(config.logging.extra_loggers), }, ], "name": "flow-container", "image": "test/image:tag", "command": [ "/bin/sh", "-c", "python -c 'import prefect; prefect.environments.execution.load_and_run_flow()'", ], }], } boto3_client = MagicMock() boto3_client.describe_task_definition.return_value = { "taskDefinition": existing_task_definition } monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) environment = FargateTaskEnvironment() environment.setup( Flow( "test", storage=Docker(registry_url="test", image_name="image", image_tag="tag"), )) assert boto3_client.describe_task_definition.called assert not boto3_client.register_task_definition.called
def test_get_flow_image_docker_storage(): flow = Flow( "test", environment=LocalEnvironment(), storage=Docker(registry_url="test", image_name="name", image_tag="tag"), ) image = get_flow_image(flow=flow) assert image == "test/name:tag"
def test_extra_dockerfile_commands(): with tempfile.TemporaryDirectory() as directory: storage = Docker( extra_dockerfile_commands=[ 'RUN echo "I\'m a little tea pot"', ], ) storage.add_flow(Flow("foo")) dpath = storage.create_dockerfile_object(directory=directory) with open(dpath, "r") as dockerfile: output = dockerfile.read() assert "COPY {} /path/test_file.txt".format( 'RUN echo "I\'m a little tea pot"\n' in output ), output
def test_setup_definition_register_no_defintions(monkeypatch): boto3_client = MagicMock() boto3_client.describe_task_definition.side_effect = ClientError({}, None) boto3_client.register_task_definition.return_value = {} monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client)) environment = FargateTaskEnvironment(family="test") environment.setup( Flow( "test", storage=Docker(registry_url="test", image_name="image", image_tag="tag"), )) assert boto3_client.describe_task_definition.called assert boto3_client.register_task_definition.called assert boto3_client.register_task_definition.call_args[1][ "family"] == "test" assert boto3_client.register_task_definition.call_args[1][ "containerDefinitions"] == [{ "environment": [ { "name": "PREFECT__CLOUD__GRAPHQL", "value": prefect.config.cloud.graphql, }, { "name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false" }, { "name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS", "value": "prefect.engine.cloud.CloudFlowRunner", }, { "name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS", "value": "prefect.engine.cloud.CloudTaskRunner", }, { "name": "PREFECT__CLOUD__SEND_FLOW_RUN_LOGS", "value": "true" }, { "name": "PREFECT__LOGGING__EXTRA_LOGGERS", "value": "[]", }, ], "name": "flow-container", "image": "test/image:tag", "command": [ "/bin/sh", "-c", "python -c 'import prefect; prefect.environments.execution.load_and_run_flow()'", ], }]
def test_get_flow_image_docker_storage(): flow = Flow( "test", run_config=UniversalRun(), storage=Docker(registry_url="test", image_name="name", image_tag="tag"), ) image = get_flow_image(flow=flow) assert image == "test/name:tag"
def test_create_dockerfile_from_dockerfile(): myfile = "FROM my-own-image:latest\n\nRUN echo 'hi'" with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "myfile"), "w") as tmp: tmp.write(myfile) storage = Docker(dockerfile=os.path.join(directory, "myfile")) dpath = storage.create_dockerfile_object(directory=directory) with open(dpath, "r") as dockerfile: output = dockerfile.read() assert output.startswith("\n" + myfile) # test proper indentation assert all(line == line.lstrip() for line in output.split("\n") if line not in ["\n", " "])
def test_create_namespaced_job_fails_outside_cluster(job_spec_file): environment = KubernetesJobEnvironment(job_spec_file=job_spec_file) storage = Docker(registry_url="test1", image_name="test2", image_tag="test3") with pytest.raises(EnvironmentError): with set_temporary_config({"cloud.auth_token": "test"}): environment.execute(Flow("test", storage=storage))
def test_deploy_flow_errors_if_mix_task_definition_arn_and_docker_storage(self): with pytest.raises( ValueError, match="Cannot provide `task_definition_arn` when using `Docker` storage", ): self.deploy_flow( ECSRun(task_definition_arn="my-taskdef-arn"), storage=Docker(registry_url="test", image_name="name", image_tag="tag"), )
def test_build_with_rm_override(monkeypatch): storage = Docker( registry_url="reg", base_image="python:3.7", image_name="test", image_tag="latest", build_kwargs={"rm": False}, ) mock_docker_client = MagicMock() mock_docker_client.images.return_value = ["test"] with patch.object(storage, "_get_client") as mock_docker_client_fn: mock_docker_client_fn.return_value = mock_docker_client output = storage.build(push=False) mock_docker_client.build.assert_called_once_with( dockerfile=ANY, path=ANY, tag="reg/test:latest", rm=False )
def test_run_healthchecks_arg_custom_prefect_dir(ignore_healthchecks, tmpdir): with open(os.path.join(tmpdir, "test"), "w+") as t: t.write("asdf") storage = Docker(ignore_healthchecks=ignore_healthchecks, prefect_directory="/usr/local/prefect") f = Flow("test") storage.add_flow(f) dpath = storage.create_dockerfile_object(directory=tmpdir) with open(dpath, "r") as dockerfile: output = dockerfile.read() if ignore_healthchecks: assert "RUN python /usr/local/prefect/healthcheck.py" not in output else: assert "RUN python /usr/local/prefect/healthcheck.py" in output
def test_create_dockerfile_with_weird_flow_name(): with tempfile.TemporaryDirectory() as tempdir_outside: with open(os.path.join(tempdir_outside, "test"), "w+") as t: t.write("asdf") with tempfile.TemporaryDirectory() as tempdir: storage = Docker(registry_url="test1", base_image="test3") f = Flow("WHAT IS THIS !!! ~~~~") storage.add_flow(f) dpath = storage.create_dockerfile_object(directory=tempdir) with open(dpath, "r") as dockerfile: output = dockerfile.read() assert ( "COPY what-is-this.flow /opt/prefect/flows/what-is-this.prefect" in output)
def configure_docker(): # Using Docker fp = Path(__file__) flow.storage = Docker( python_dependencies=["git+https://github.com/steph-ben/datafetch.git"], stored_as_script=True, path=f"/flow/{ fp.name }", files={fp.absolute(): f"/flow/{ fp.name }"}, build_kwargs={'nocache': False}) flow.run_config = DockerRun()
def test_create_dockerfile_with_weird_flow_name_custom_prefect_dir(tmpdir): with open(os.path.join(tmpdir, "test"), "w+") as t: t.write("asdf") storage = Docker( registry_url="test1", base_image="test3", prefect_directory="/tmp/prefect-is-c00l", ) f = Flow("WHAT IS THIS !!! ~~~~") storage.add_flow(f) dpath = storage.create_dockerfile_object(directory=tmpdir) with open(dpath, "r") as dockerfile: output = dockerfile.read() assert ( "COPY what-is-this.flow /tmp/prefect-is-c00l/flows/what-is-this.prefect" in output)
def test_build_with_default_rm_true(monkeypatch): storage = Docker( registry_url="reg", base_image="python:3.7", image_name="test", image_tag="latest", ) pull_image = MagicMock() monkeypatch.setattr("prefect.storage.Docker.pull_image", pull_image) mock_docker_client = MagicMock() mock_docker_client.images.return_value = ["test"] with patch.object(storage, "_get_client") as mock_docker_client_fn: mock_docker_client_fn.return_value = mock_docker_client output = storage.build(push=False) mock_docker_client.build.assert_called_once_with( dockerfile=ANY, path=ANY, tag="reg/test:latest", rm=True )
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_create_namespaced_job_fails_outside_cluster(): environment = DaskKubernetesEnvironment() storage = Docker(registry_url="test1", image_name="test2", image_tag="test3") with pytest.raises(EnvironmentError): with set_temporary_config({"cloud.auth_token": "test"}): flow = base_flow flow.storage = storage with set_temporary_config({"cloud.auth_token": "test"}): environment.execute(flow=flow)
def test_run_healthchecks_arg(ignore_healthchecks): with tempfile.TemporaryDirectory() as tempdir_outside: with open(os.path.join(tempdir_outside, "test"), "w+") as t: t.write("asdf") with tempfile.TemporaryDirectory() as tempdir: storage = Docker(ignore_healthchecks=ignore_healthchecks) f = Flow("test") storage.add_flow(f) dpath = storage.create_dockerfile_object(directory=tempdir) with open(dpath, "r") as dockerfile: output = dockerfile.read() if ignore_healthchecks: assert "RUN python /opt/prefect/healthcheck.py" not in output else: assert "RUN python /opt/prefect/healthcheck.py" in output
def test_docker_storage_get_flow_method(): with tempfile.TemporaryDirectory() as directory: storage = Docker(base_image="python:3.6", prefect_directory=directory) with pytest.raises(ValueError): storage.get_flow() @prefect.task def add_to_dict(): with open(os.path.join(directory, "output"), "w") as tmp: tmp.write("success") flow_dir = os.path.join(directory, "flows") os.makedirs(flow_dir, exist_ok=True) with open(os.path.join(flow_dir, "test.prefect"), "w+") as env: flow = Flow("test", tasks=[add_to_dict]) flow_path = os.path.join(flow_dir, "test.prefect") with open(flow_path, "wb") as f: cloudpickle.dump(flow, f) out = storage.add_flow(flow) f = storage.get_flow(out) assert isinstance(f, Flow) assert f.name == "test" assert len(f.tasks) == 1
def test_copy_files_with_dockerignore(): with tempfile.TemporaryDirectory() as sample_top_directory: sample_sub_directory = os.path.join(sample_top_directory, "subdir") os.mkdir(sample_sub_directory) sample_file = os.path.join(sample_sub_directory, "test.txt") with open(sample_file, "w+") as t: t.write("asdf") dockerignore = os.path.join(sample_sub_directory, ".dockerignore") with open(dockerignore, "w+") as t: t.write("test.txt") with tempfile.TemporaryDirectory() as directory: storage = Docker( files={ sample_sub_directory: "/test_dir", sample_file: "/path/test_file.txt", }, dockerignore=dockerignore, ) storage.add_flow(Flow("foo")) storage.create_dockerfile_object(directory=directory) contents = os.listdir(directory) assert ".dockerignore" in contents, contents
def test_dockerfile_env_vars(tmpdir): env_vars = OrderedDict([ ("NUM", 1), ("STR_WITH_SPACES", "Hello world!"), ("STR_WITH_QUOTES", 'Hello "friend"'), ("STR_WITH_SINGLE_QUOTES", "'foo'"), ]) storage = Docker(env_vars=env_vars, ) storage.add_flow(Flow("foo")) dpath = storage.create_dockerfile_object(directory=str(tmpdir)) with open(dpath, "r") as dockerfile: output = dockerfile.read() expected = textwrap.dedent(""" ENV NUM=1 \\ STR_WITH_SPACES='Hello world!' \\ STR_WITH_QUOTES='Hello "friend"' \\ STR_WITH_SINGLE_QUOTES="'foo'" \\ """) assert expected in output
def test_build_sets_image_name_for_multiple_flows(monkeypatch): storage = Docker(registry_url="reg") storage.add_flow(Flow("test")) storage.add_flow(Flow("test2")) monkeypatch.setattr("prefect.storage.Docker._build_image", MagicMock()) output = storage.build() assert output.registry_url == storage.registry_url assert isinstance(output.image_name, str) assert output.image_tag.startswith(str(pendulum.now("utc").year))
def test_build_image_passes(monkeypatch): flow = Flow("test") storage = Docker( registry_url="reg", base_image="python:3.6", image_name="test", image_tag="latest", ) pull_image = MagicMock() monkeypatch.setattr("prefect.storage.Docker.pull_image", pull_image) build = MagicMock() monkeypatch.setattr("docker.APIClient.build", build) images = MagicMock(return_value=["test"]) monkeypatch.setattr("docker.APIClient.images", images) image_name, image_tag = storage._build_image(flow, push=False) assert image_name assert image_tag
def test_build_respects_user_provided_image_name_and_tag_for_multiple_flows( monkeypatch, ): storage = Docker(registry_url="reg", image_name="CUSTOM", image_tag="TAG") storage.add_flow(Flow("test")) storage.add_flow(Flow("test2")) monkeypatch.setattr("prefect.storage.Docker._build_image", MagicMock()) output = storage.build() assert output.registry_url == storage.registry_url assert output.image_name == "CUSTOM" assert output.image_tag == "TAG"