def test_copy_files():
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")
with tempfile.TemporaryDirectory() as directory:
storage = Docker(
files={
sample_sub_directory: "/test_dir",
sample_file: "/path/test_file.txt",
},
)
storage.add_flow(Flow("foo"))
dpath = storage.create_dockerfile_object(directory=directory)
with open(dpath, "r") as dockerfile:
output = dockerfile.read()
contents = os.listdir(directory)
assert "subdir" in contents, contents
assert "test.txt" in contents, contents
assert "COPY {} /test_dir".format(
os.path.join(directory, "subdir").replace("\\", "/") in output
), output
assert "COPY {} /path/test_file.txt".format(
os.path.join(directory, "test.txt").replace("\\", "/") in output
), output
def test_create_dockerfile_from_everything(no_docker_host_var):
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",
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
assert "COPY healthcheck.py /opt/prefect/healthcheck.py" in output
assert "COPY test.flow /opt/prefect/flows/test.prefect" in output
assert "COPY other.flow /opt/prefect/flows/other.prefect" 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_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_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_build_sets_informative_image_name_for_weird_name_flows(monkeypatch):
storage = Docker(registry_url="reg")
storage.add_flow(Flow("!&& ~~ cool flow :shades:"))
monkeypatch.setattr("prefect.storage.Docker._build_image", MagicMock())
output = storage.build()
assert output.registry_url == storage.registry_url
assert output.image_name == "cool-flow-shades"
assert output.image_tag.startswith(str(pendulum.now("utc").year))
def test_build_respects_user_provided_image_name_and_tag(monkeypatch):
storage = Docker(registry_url="reg", image_name="CUSTOM", image_tag="TAG")
storage.add_flow(Flow("test"))
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"
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_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(tmpdir):
flow_dir = str(tmpdir.mkdir("flows"))
flow = Flow("test")
flow_path = os.path.join(flow_dir, "test.prefect")
with open(flow_path, "wb") as f:
cloudpickle.dump(flow, f)
storage = Docker(base_image="python:3.6", prefect_directory=str(tmpdir))
storage.add_flow(flow)
f = storage.get_flow(flow.name)
assert isinstance(f, Flow)
assert f.name == "test"
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) == "/opt/prefect/flows/test.prefect"
assert f.name in storage
assert storage.flows[f.name] == "/opt/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_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 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_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_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 test_add_similar_flows_fails():
storage = Docker()
flow = prefect.Flow("test")
storage.add_flow(flow)
with pytest.raises(ValueError):
storage.add_flow(flow)
