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.environments.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"
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_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_add_flow_to_docker(): storage = Docker() f = Flow("test") assert f not in storage assert storage.add_flow(f) == "/root/.prefect/flows/test.prefect" assert f.name in storage assert storage.flows[f.name] == "/root/.prefect/flows/test.prefect"
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_add_flow_with_weird_name_is_cleaned(): storage = Docker() flow = prefect.Flow("WELL what do you know?!~? looks like a test!!!!") loc = storage.add_flow(flow) assert "?" not in loc assert "!" not in loc assert " " not in loc assert "~" not in loc
def test_add_flow_to_docker_custom_prefect_dir(): storage = Docker(prefect_directory="/usr/prefect-is-gr8") f = Flow("test") assert f not in storage assert storage.add_flow(f) == "/usr/prefect-is-gr8/flows/test.prefect" assert f.name in storage assert storage.flows[f.name] == "/usr/prefect-is-gr8/flows/test.prefect" assert storage.env_vars == { "PREFECT__USER_CONFIG_PATH": "/usr/prefect-is-gr8/config.toml" }
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 /root/.prefect/flows/what-is-this.prefect" in output)
def test_create_dockerfile_from_everything(): 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", python_dependencies=["test"], image_name="test4", image_tag="test5", env_vars={"test": "1"}, files={os.path.join(tempdir_outside, "test"): "./test2"}, base_url="test_url", ) f = Flow("test") g = Flow("other") storage.add_flow(f) storage.add_flow(g) dpath = storage.create_dockerfile_object(directory=tempdir) with open(dpath, "r") as dockerfile: output = dockerfile.read() assert "FROM test3" in output assert "COPY test ./test2" in output # ensure there is a "ENV ... test=1" in the output results = re.search( r"ENV(\s+[a-zA-Z0-9_]*\=[^\\]*\\\s*$)*\s*(?P<result>test=1)", output, re.MULTILINE, ) assert results != None assert results.group("result") == "test=1" assert "COPY healthcheck.py /root/.prefect/healthcheck.py" in output assert "COPY test.flow /root/.prefect/flows/test.prefect" in output assert "COPY other.flow /root/.prefect/flows/other.prefect" in output
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 create_flow(): """ Create the flow """ result = S3Result(bucket="my-prefect-flows", ) with Flow("Sample Flow", result=result) as flow: first_result = add(1, y=2) second_result = add(x=first_result, y=100) storage = Docker(image_name="asdf", image_tag="no", python_dependencies=["boto3"], secrets=["AWS_CREDENTIALS"], prefect_version="aws_creds_fix") storage.add_flow(flow) flow.storage = storage return flow
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_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 create_flow(): """ Create the flow """ result = S3Result(bucket="my-prefect-flows", ) with Flow("Sample Flow #1", result=result) as flow: first_result = add(1, y=2) second_result = add(x=first_result, y=100) storage = Docker( base_image="prefecthq/prefect:0.11.5-python3.7", image_name="sample-flow", image_tag="test2", python_dependencies=["boto3"], prefect_version="master", build_kwargs=dict(), ) storage.add_flow(flow) flow.storage = storage return flow
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_create_dockerfile_with_flow_file(no_docker_host_var, tmpdir): contents = """from prefect import Flow\nf=Flow('test-flow')""" full_path = os.path.join(tmpdir, "flow.py") with open(full_path, "w") as f: f.write(contents) with open(os.path.join(tmpdir, "test"), "w+") as t: t.write("asdf") with tempfile.TemporaryDirectory() as tempdir_inside: storage = Docker( files={full_path: "flow.py"}, stored_as_script=True, path="flow.py", ) f = Flow("test-flow") storage.add_flow(f) dpath = storage.create_dockerfile_object(directory=tempdir_inside) with open(dpath, "r") as dockerfile: output = dockerfile.read() assert "COPY flow.py flow.py" in output storage = Docker( files={full_path: "flow.py"}, stored_as_script=True, ) f = Flow("test-flow") storage.add_flow(f) with pytest.raises(ValueError): storage.create_dockerfile_object(directory=tempdir_inside)
def test_add_similar_flows_fails(): storage = Docker() flow = prefect.Flow("test") storage.add_flow(flow) with pytest.raises(ValueError): storage.add_flow(flow)
from myflows.thisflow import flow as tflow1 # from myflows.thisflow import flow as tflow2 # tflow.storage.build() from prefect.environments.storage import Docker import copy tflow2 = copy.deepcopy(tflow1) tflow2.name = "222" d = Docker() d.add_flow(tflow1) d.add_flow(tflow2) d.build()
@task def numbers_task(): return [1, 2, 3] @task def map_task(x): return x + 1 @task def reduce_task(x): return sum(x) with Flow("Map / Reduce 🤓") as mr_flow: numbers = numbers_task() first_map = map_task.map(numbers) second_map = map_task.map(first_map) reduction = reduce_task(second_map) storage.add_flow(etl_flow) storage.add_flow(mr_flow) storage = storage.build(push=False) etl_flow.storage = storage mr_flow.storage = storage etl_flow.deploy(build="False") mr_flow.deploy(build="False")