def test_add_flow_raises_if_name_conflict(tmpdir):
storage = Local(directory=str(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_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_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_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_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():
with tempfile.TemporaryDirectory() as tmpdir:
s = Local(directory=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_containment(tmpdir):
s = Local(directory=str(tmpdir))
f = Flow("test")
s.add_flow(f)
assert True not in s
assert f not in s
assert "test" in s
assert Flow("other") not in s
assert "other" not in s
def test_containment():
with tempfile.TemporaryDirectory() as tmpdir:
s = Local(directory=tmpdir)
f = Flow("test")
s.add_flow(f)
assert True not in s
assert f not in s
assert "test" in s
assert Flow("other") not in s
assert "other" not in s
def test_get_flow_from_file_returns_flow(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)
f = Flow("test-flow")
storage = Local(stored_as_script=True, path=full_path)
storage.add_flow(f)
flow = storage.get_flow(full_path)
assert flow.run()
def build_flow_run(self, config, storage=None, core_version="0.13.0"):
if storage is None:
storage = Local()
return GraphQLResult({
"flow":
GraphQLResult({
"storage": storage.serialize(),
"id": "new_id",
"core_version": core_version,
}),
"run_config":
None if config is None else config.serialize(),
"id":
"id",
})
def test_add_flow_to_storage():
with tempfile.TemporaryDirectory() as tmpdir:
storage = Local(directory=tmpdir)
f = Flow("test")
assert f.name not in storage
res = storage.add_flow(f)
assert res.endswith("test.prefect")
assert f.name in storage
with open(os.path.join(tmpdir, "test.prefect"), "rb") as f:
wat = f.read()
assert isinstance(wat, bytes)
assert cloudpickle.loads(wat).name == "test"
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_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_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_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 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 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_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 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_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
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_multiple_flows_in_storage(tmpdir):
s = Local(directory=str(tmpdir))
f = Flow("test")
g = Flow("other")
f_loc = s.add_flow(f)
g_loc = s.add_flow(g)
assert "test" in s
assert "other" in s
assert "not" not in s
assert s.get_flow(f.name) == f
assert s.get_flow(g.name) == g
assert s.flows["test"] == f_loc
assert s.flows["other"] == g_loc
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
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
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
def test_add_flow_to_storage(tmpdir):
storage = Local(directory=str(tmpdir))
f = Flow("test")
assert f.name not in storage
res = storage.add_flow(f)
flow_dir = os.path.join(tmpdir, "test")
assert os.path.exists(flow_dir)
assert len(os.listdir(flow_dir)) == 1
assert res.startswith(flow_dir)
assert f.name in storage
f2 = storage.get_flow(f.name)
assert f2.name == "test"
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_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