def test_get_flow_returns_flow():
with tempfile.TemporaryDirectory() as tmpdir:
storage = Local(directory=tmpdir)
f = Flow("test")
loc = storage.add_flow(f)
runner = storage.get_flow(loc)
assert runner == f
def test_k8s_agent_deploy_flow_raises(monkeypatch, cloud_api):
batch_client = MagicMock()
monkeypatch.setattr("kubernetes.client.BatchV1Api",
MagicMock(return_value=batch_client))
core_client = MagicMock()
core_client.list_namespaced_pod.return_value = MagicMock(items=[])
monkeypatch.setattr("kubernetes.client.CoreV1Api",
MagicMock(return_value=core_client))
get_jobs = MagicMock(return_value=[])
monkeypatch.setattr(
"prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
get_jobs,
)
agent = KubernetesAgent()
with pytest.raises(ValueError):
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"id": "id",
"environment": LocalEnvironment().serialize(),
"core_version": "0.13.0",
}),
"id":
"id",
}))
assert not agent.batch_client.create_namespaced_job.called
async def flow(self, project_id, tmpdir):
"""
A simple diamond flow
"""
@prefect.task
def numbers():
return [1, 2, 3]
@prefect.task
def add(x, y):
return x + y
with prefect.Flow(
"simple map",
run_config=prefect.run_configs.LocalRun(),
storage=Local(directory=tmpdir),
) as flow:
flow.numbers1 = numbers()
flow.numbers2 = numbers()
flow.add = add.map(flow.numbers1, flow.numbers2)
add_state_handler(flow)
with set_temporary_config(key="dev", value=True):
flow.server_id = await api.flows.create_flow(
project_id=project_id,
serialized_flow=flow.serialize(build=True))
return flow
async def flow(self, project_id, tmpdir):
"""
A simple diamond flow
"""
@prefect.task
def numbers():
return [1, 2, 3]
@prefect.task
def add(x):
if x == 2:
raise prefect.engine.signals.FAIL("Don't run for 2")
elif not isinstance(x, int):
return -99
return x + 10
with prefect.Flow(
"simple map",
run_config=prefect.run_configs.LocalRun(),
storage=Local(directory=tmpdir),
) as flow:
flow.numbers = numbers()
flow.add1 = add.map(flow.numbers)
flow.add2 = add.map(flow.add1)
flow.add2.trigger = prefect.triggers.all_finished
add_state_handler(flow)
with set_temporary_config(key="dev", value=True):
flow.server_id = await api.flows.create_flow(
project_id=project_id,
serialized_flow=flow.serialize(build=True))
return flow
async def flow(self, project_id, tmpdir):
"""
A simple diamond flow
"""
@prefect.task
def numbers():
return [1, 2, 3]
@prefect.task
def add(x):
return x + 1
@prefect.task
def get_sum(x):
return sum(x)
with prefect.Flow(
"simple reduce",
environment=LocalEnvironment(),
storage=Local(directory=tmpdir),
) as flow:
flow.numbers = numbers()
flow.add = add.map(flow.numbers)
flow.sum = get_sum(flow.add)
add_state_handler(flow)
with set_temporary_config(key="dev", value=True):
flow.server_id = await api.flows.create_flow(
project_id=project_id,
serialized_flow=flow.serialize(build=True))
return flow
def test_local_agent_deploy_run_config_missing_working_dir(monkeypatch, tmpdir):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
working_dir = str(tmpdir.join("missing"))
agent = LocalAgent()
with pytest.raises(ValueError, match="nonexistent `working_dir`"):
agent.deploy_flow(
flow_run=GraphQLResult(
{
"id": "id",
"flow": {
"storage": Local().serialize(),
"id": "foo",
"core_version": "0.13.0",
},
"run_config": LocalRun(working_dir=working_dir).serialize(),
},
)
)
assert not popen.called
assert not agent.processes
async def flow(self, project_id, tmpdir):
"""
A simple diamond flow
"""
@prefect.task
def numbers():
return [1, 2, 3]
@prefect.task
def add(x):
if x == 2:
raise prefect.engine.signals.FAIL("Don't run for 2")
return x + 10
with prefect.Flow(
"simple map",
environment=LocalEnvironment(),
storage=Local(directory=tmpdir),
) as flow:
flow.numbers = numbers()
flow.add1 = add.map(flow.numbers)
flow.add2 = add.map(flow.add1)
add_state_handler(flow)
with set_temporary_config(key="dev", value=True):
flow.server_id = await api.flows.create_flow(
project_id=project_id,
serialized_flow=flow.serialize(build=True))
return flow
def test_local_agent_deploy_run_config_working_dir(monkeypatch, working_dir, tmpdir):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
if working_dir is not None:
working_dir = str(tmpdir)
agent = LocalAgent()
agent.deploy_flow(
flow_run=GraphQLResult(
{
"id": "id",
"flow": {
"storage": Local().serialize(),
"id": "foo",
"core_version": "0.13.0",
},
"run_config": LocalRun(working_dir=working_dir).serialize(),
},
)
)
assert popen.called
assert len(agent.processes) == 1
assert popen.call_args[1]["cwd"] == working_dir
def test_local_agent_deploy_unsupported_run_config(monkeypatch):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent()
with pytest.raises(
TypeError,
match="`run_config` of type `KubernetesRun`, only `LocalRun` is supported",
):
agent.deploy_flow(
flow_run=GraphQLResult(
{
"id": "id",
"flow": {
"storage": Local().serialize(),
"id": "foo",
"core_version": "0.13.0",
},
"run_config": KubernetesRun().serialize(),
},
)
)
assert not popen.called
assert len(agent.processes) == 0
def configure_local():
# Using Local
flow.storage = Local(
stored_as_script=True,
path=
"/home/steph/Code/steph-ben/fetch-with-prefect/config/laptop/simpleflow.py"
)
def test_docker_agent_deploy_flow_does_not_include_host_gateway_for_old_engine_versions(
api, docker_engine_version):
api.version.return_value = {"Version": docker_engine_version}
run = UniversalRun()
storage = Local()
agent = DockerAgent()
with pytest.warns(
UserWarning,
match=
("`host.docker.internal` could not be automatically resolved.*"
f"feature is not supported on Docker Engine v{docker_engine_version}"
),
):
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 "extra_hosts" not in api.create_host_config.call_args[1]
def test_add_flow_raises_if_name_conflict(tmpdir):
storage = Local(directory=str(tmpdir))
f = Flow("test")
storage.add_flow(f)
g = Flow("test")
with pytest.raises(ValueError, match='name "test"'):
storage.add_flow(g)
def test_build_returns_self(tmpdir):
s = Local(directory=str(tmpdir))
assert s.build() is s
f = Flow("test")
s.add_flow(f)
assert s.build() is s
def load_flows_from_script(path: str) -> "List[prefect.Flow]":
"""Given a file path, load all flows found in the file"""
# We use abs_path for everything but logging (logging the original
# user-specified path provides a clearer message).
abs_path = os.path.abspath(path)
# Temporarily add the flow's local directory to `sys.path` so that local
# imports work. This ensures that `sys.path` is the same as it would be if
# the flow script was run directly (i.e. `python path/to/flow.py`).
orig_sys_path = sys.path.copy()
sys.path.insert(0, os.path.dirname(abs_path))
try:
with prefect.context({
"loading_flow": True,
"local_script_path": abs_path
}):
namespace = runpy.run_path(abs_path, run_name="<flow>")
except Exception as exc:
click.secho(f"Error loading {path!r}:", fg="red")
log_exception(exc, 2)
raise TerminalError
finally:
sys.path[:] = orig_sys_path
flows = [f for f in namespace.values() if isinstance(f, prefect.Flow)]
if flows:
for f in flows:
if f.storage is None:
f.storage = Local(path=abs_path, stored_as_script=True)
return flows
def test_docker_agent_deploy_flow_storage_raises(monkeypatch, api):
monkeypatch.setattr("prefect.agent.agent.Client", MagicMock())
agent = DockerAgent()
with pytest.raises(ValueError):
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Local().serialize(),
"id": "foo",
"environment": LocalEnvironment().serialize(),
"core_version": "0.13.0",
}
),
"id": "id",
"name": "name",
"version": "version",
}
)
)
assert not api.pull.called
def test_docker_agent_deploy_flow_unsupported_run_config(api):
agent = DockerAgent()
with pytest.raises(
TypeError,
match=
"`run_config` of type `LocalRun`, only `DockerRun` is supported",
):
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"id": "foo",
"name": "flow-name",
"core_version": "0.13.0",
}),
"run_config":
LocalRun().serialize(),
"id":
"id",
"name":
"name",
"version":
"version",
}))
assert not api.pull.called
def test_docker_agent_deploy_flow_no_pull_using_environment_metadata(api):
agent = DockerAgent(no_pull=True)
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"name":
"flow-name",
"storage":
Local().serialize(),
"environment":
LocalEnvironment(metadata={
"image": "name:tag"
}).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_docker_agent_deploy_flow_sets_container_name_with_slugify(
api, run_name, container_name):
"""
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
"""
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":
run_name,
}))
assert api.create_container.call_args[1]["name"] == container_name
def get_run_task_kwargs(self,
run_config,
tenant_id: str = None,
taskdef=None,
**kwargs):
agent = ECSAgent(**kwargs)
if tenant_id:
agent.client._get_auth_tenant = MagicMock(return_value=tenant_id)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"id": "flow-id",
"version": 1,
"name": "Test Flow",
"core_version": "0.13.0",
}),
"run_config":
run_config.serialize(),
"id":
"flow-run-id",
})
if taskdef is None:
taskdef = agent.generate_task_definition(flow_run, run_config)
return agent.get_run_task_kwargs(flow_run, run_config, taskdef)
def test_docker_agent_deploy_flow_uses_environment_metadata(api):
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":
"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_get_flow_image_raises_on_missing_info():
flow = Flow(
"test",
environment=LocalEnvironment(),
storage=Local(),
)
with pytest.raises(ValueError):
get_flow_image(flow=flow)
def test_get_flow_image_raises_on_missing_info():
flow = Flow(
"test",
run_config=UniversalRun(),
storage=Local(),
)
with pytest.raises(ValueError):
get_flow_image(flow=flow)
def test_add_flow_raises_if_name_conflict():
with tempfile.TemporaryDirectory() as tmpdir:
storage = Local(directory=tmpdir)
f = Flow("test")
res = storage.add_flow(f)
g = Flow("test")
with pytest.raises(ValueError, match='name "test"'):
storage.add_flow(g)
def test_add_flow_with_weird_name_is_cleaned(tmpdir):
s = Local(directory=str(tmpdir))
flow = Flow("WELL what do you know?!~? looks like a test!!!!")
loc = s.add_flow(flow)
assert "?" not in loc
assert "!" not in loc
assert " " not in loc
assert "~" not in loc
def test_create_local_storage_without_validation():
storage = Local(directory="C:\\Users\\chris\\.prefect\\flows",
validate=False)
assert storage
assert storage.directory == "C:\\Users\\chris\\.prefect\\flows"
assert isinstance(storage.result, LocalResult)
assert storage.result.dir == storage.directory
def test_build_returns_self():
with tempfile.TemporaryDirectory() as tmpdir:
s = Local(directory=tmpdir)
assert s.build() is s
f = Flow("test")
s.add_flow(f)
assert s.build() is s
def test_get_flow_image_env_metadata():
flow = Flow(
"test",
environment=LocalEnvironment(metadata={"image": "repo/name:tag"}),
storage=Local(),
)
image = get_flow_image(flow=flow)
assert image == "repo/name:tag"
def test_get_flow_image_run_config():
flow = Flow(
"test",
run_config=DockerRun(image="repo/name:tag"),
storage=Local(),
)
image = get_flow_image(flow=flow)
assert image == "repo/name:tag"
def test_client_register_flow_id_output(
patch_post, use_run_config, compressed, monkeypatch, capsys, cloud_api, tmpdir
):
if compressed:
response = {
"data": {
"project": [{"id": "proj-id"}],
"create_flow_from_compressed_string": {"id": "long-id"},
"flow_by_pk": {"flow_group_id": "fg-id"},
}
}
else:
response = {
"data": {
"project": [{"id": "proj-id"}],
"create_flow": {"id": "long-id"},
"flow_by_pk": {"flow_group_id": "fg-id"},
}
}
patch_post(response)
monkeypatch.setattr(
"prefect.client.Client.get_default_tenant_slug", MagicMock(return_value="tslug")
)
with set_temporary_config(
{
"cloud.api": "http://my-cloud.foo",
"cloud.auth_token": "secret_token",
"backend": "cloud",
}
):
client = Client()
labels = ["test1", "test2"]
storage = Local(tmpdir)
if use_run_config:
flow = prefect.Flow(
name="test", storage=storage, run_config=LocalRun(labels=labels)
)
flow.environment = None
else:
flow = prefect.Flow(
name="test", storage=storage, environment=LocalEnvironment(labels=labels)
)
flow.result = flow.storage.result
flow_id = client.register(
flow,
project_name="my-default-project",
compressed=compressed,
version_group_id=str(uuid.uuid4()),
)
assert flow_id == "long-id"
captured = capsys.readouterr()
assert "Flow URL: https://cloud.prefect.io/tslug/flow/fg-id\n" in captured.out
assert f"Labels: {labels}" in captured.out
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"}
host_config = {"auto_remove": False, "shm_size": "128m"}
exp_host_config = {
"auto_remove": False,
"extra_hosts": {
"host.docker.internal": "host-gateway"
},
"shm_size": "128m",
}
run = DockerRun(image=image, env=env, host_config=host_config)
else:
env = {}
host_config = {}
exp_host_config = {
"auto_remove": True,
"extra_hosts": {
"host.docker.internal": "host-gateway"
},
}
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
res_host_config = api.create_host_config.call_args[1]
for k, v in exp_host_config.items():
assert res_host_config[k] == v