def test_get_flow_image_raises_on_missing_info():
flow_run = GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Local().serialize(),
"environment": LocalEnvironment().serialize(),
"id": "id",
}
),
"id": "id",
}
)
with pytest.raises(ValueError):
image = get_flow_image(flow_run=flow_run)
def test_local_agent_deploy_import_paths(monkeypatch, runner_token):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent(import_paths=["paths"])
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({"storage": Local(directory="test").serialize()}),
"id":
"id",
}))
assert popen.called
assert len(agent.processes) == 1
assert "paths" in popen.call_args[1]["env"]["PYTHONPATH"]
def test_local_agent_deploy_processes_local_storage(monkeypatch, runner_token):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent()
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({"storage": Local(directory="test").serialize()}),
"id":
"id",
}))
assert popen.called
assert len(agent.processes) == 1
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 test_get_flow_image_run_config_default_value_from_core_version(version):
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"core_version": version,
"storage": Local().serialize(),
"run_config": KubernetesRun().serialize(),
"id": "id",
}),
"id":
"id",
})
image = get_flow_image(flow_run)
expected_version = version.split("+")[0] if version else "latest"
assert image == f"prefecthq/prefect:all_extras-{expected_version}"
def get_run_task_kwargs(self, run_config, **kwargs):
agent = ECSAgent(**kwargs)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"run_config": run_config.serialize(),
"id": "flow-id",
"version": 1,
"name": "Test Flow",
"core_version": "0.13.0",
}),
"id":
"flow-run-id",
})
return agent.get_run_task_kwargs(flow_run, run_config)
def test_get_flow_run_command(core_version, command):
flow_run = GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Local().serialize(),
"environment": LocalEnvironment().serialize(),
"id": "id",
"core_version": core_version,
}
),
"id": "id",
}
)
assert get_flow_run_command(flow_run) == command
def test_get_flow_run_command_works_if_core_version_not_on_response():
legacy_command = "prefect execute cloud-flow"
flow_run = GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Local().serialize(),
"environment": LocalEnvironment().serialize(),
"id": "id",
}
),
"id": "id",
}
)
assert get_flow_run_command(flow_run) == legacy_command
def test_get_flow_image_run_config_image_on_RunConfig():
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage":
Local().serialize(),
"run_config":
KubernetesRun(image="myfancyimage").serialize(),
"id":
"id",
}),
"id":
"id",
})
image = get_flow_image(flow_run)
assert image == "myfancyimage"
def test_k8s_agent_deploy_flow_uses_environment_metadata(monkeypatch, cloud_api):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
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()
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Local().serialize(),
"environment": LocalEnvironment(
metadata={"image": "repo/name:tag"}
).serialize(),
"id": "id",
"core_version": "0.13.0",
}
),
"id": "id",
}
)
)
assert agent.batch_client.create_namespaced_job.called
assert (
agent.batch_client.create_namespaced_job.call_args[1]["body"]["spec"][
"template"
]["spec"]["containers"][0]["image"]
== "repo/name:tag"
)
def test_run_flow_calls_callbacks(monkeypatch):
start_func = MagicMock()
exit_func = MagicMock()
environment = DaskKubernetesEnvironment(on_start=start_func, on_exit=exit_func)
flow_runner = MagicMock()
monkeypatch.setattr(
"prefect.engine.get_default_flow_runner_class",
MagicMock(return_value=flow_runner),
)
kube_cluster = MagicMock()
monkeypatch.setattr("dask_kubernetes.KubeCluster", kube_cluster)
with tempfile.TemporaryDirectory() as directory:
d = Local(directory)
d.add_flow(prefect.Flow("name"))
gql_return = MagicMock(
return_value=MagicMock(
data=MagicMock(
flow_run=[
GraphQLResult(
{
"flow": GraphQLResult(
{"name": "name", "storage": d.serialize(),}
)
}
)
],
)
)
)
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.environments.execution.dask.k8s.Client", client)
with set_temporary_config({"cloud.auth_token": "test"}), prefect.context(
{"flow_run_id": "id"}
):
environment.run_flow()
assert flow_runner.call_args[1]["flow"].name == "name"
assert start_func.called
assert exit_func.called
def test_local_agent_deploy_flows_storage_continues(monkeypatch, runner_token):
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
monkeypatch.setattr("prefect.agent.local.agent.docker.APIClient",
MagicMock(return_value=api))
agent = LocalAgent()
agent.deploy_flows(flow_runs=[
GraphQLResult({
"flow": GraphQLResult({"storage": Local().serialize()}),
"id": "id"
})
])
assert not api.pull.called
def test_nomad_agent_deploy_flows_continues(monkeypatch, runner_token):
requests = MagicMock()
monkeypatch.setattr("prefect.agent.nomad.agent.requests", requests)
agent = NomadAgent()
agent.deploy_flows(flow_runs=[
GraphQLResult(
{
"flow": GraphQLResult({
"storage": Local().serialize(),
"id": "id"
}),
"id": "id",
})
])
assert not requests.post.called
def test_nomad_agent_deploy_flow_raises(monkeypatch, runner_token):
requests = MagicMock()
monkeypatch.setattr("prefect.agent.nomad.agent.requests", requests)
agent = NomadAgent()
with pytest.raises(ValueError):
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult({"storage": Local().serialize(), "id": "id"}),
"id": "id",
}
)
)
assert not requests.post.called
def test_get_flow_image_env_metadata():
flow_run = GraphQLResult(
{
"flow": GraphQLResult(
{
"storage": Local().serialize(),
"environment": LocalEnvironment(
metadata={"image": "repo/name:tag"}
).serialize(),
"id": "id",
}
),
"id": "id",
}
)
image = get_flow_image(flow_run=flow_run)
assert image == "repo/name:tag"
def test_multiple_flows_in_storage():
with tempfile.TemporaryDirectory() as tmpdir:
s = Local(directory=tmpdir)
f = Flow("test")
g = Flow("other")
z = Flow("not")
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_loc) == f
assert s.get_flow(g_loc) == g
assert s.flows["test"] == f_loc
assert s.flows["other"] == g_loc
def build_flow_run(self, config, storage=None):
if storage is None:
storage = Local()
return GraphQLResult({
"flow":
GraphQLResult({
"storage":
storage.serialize(),
"run_config":
None if config is None else config.serialize(),
"id":
"new_id",
"core_version":
"0.13.0",
}),
"id":
"id",
})
def generate_task_definition(self, run_config, storage=None, **kwargs):
if storage is None:
storage = Local()
agent = ECSAgent(**kwargs)
flow_run = GraphQLResult({
"flow":
GraphQLResult({
"storage": storage.serialize(),
"run_config": run_config.serialize(),
"id": "flow-id",
"version": 1,
"name": "Test Flow",
"core_version": "0.13.0",
}),
"id":
"flow-run-id",
})
return agent.generate_task_definition(flow_run, run_config)
def test_run_flow(monkeypatch):
environment = FargateTaskEnvironment(executor_kwargs={"test": "here"})
flow_runner = MagicMock()
monkeypatch.setattr(
"prefect.engine.get_default_flow_runner_class",
MagicMock(return_value=flow_runner),
)
executor = MagicMock()
monkeypatch.setattr(
"prefect.engine.get_default_executor_class", MagicMock(return_value=executor),
)
with tempfile.TemporaryDirectory() as directory:
d = Local(directory)
d.add_flow(prefect.Flow("name"))
gql_return = MagicMock(
return_value=MagicMock(
data=MagicMock(
flow_run=[
GraphQLResult(
{
"flow": GraphQLResult(
{"name": "name", "storage": d.serialize(),}
)
}
)
],
)
)
)
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.environments.execution.base.Client", client)
with set_temporary_config({"cloud.auth_token": "test"}), prefect.context(
{"flow_run_id": "id"}
):
environment.run_flow()
assert flow_runner.call_args[1]["flow"].name == "name"
assert executor.call_args[1] == {"test": "here"}
def test_load_and_run_flow(monkeypatch, tmpdir):
myflow = Flow("test-flow")
# This is gross. Since the flow is pickled/unpickled, there's no easy way
# to access the same object to set a flag. Resort to setting an environment
# variable as a global flag that won't get copied eagerly through
# cloudpickle.
monkeypatch.setenv("TEST_RUN_CALLED", "FALSE")
class MyEnvironment(Environment):
def run(self, flow):
assert flow is myflow
os.environ["TEST_RUN_CALLED"] = "TRUE"
myflow.environment = MyEnvironment()
storage = Local(str(tmpdir))
myflow.storage = storage
storage.add_flow(myflow)
gql_return = MagicMock(
return_value=MagicMock(
data=MagicMock(
flow_run=[
GraphQLResult(
{
"flow": GraphQLResult(
{"name": myflow.name, "storage": storage.serialize()}
)
}
)
],
)
)
)
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.environments.execution.base.Client", client)
with set_temporary_config({"cloud.auth_token": "test"}), prefect.context(
{"flow_run_id": "id"}
):
load_and_run_flow()
assert os.environ["TEST_RUN_CALLED"] == "TRUE"
def test_nomad_agent_deploy_flow_raises(monkeypatch, runner_token):
post = MagicMock()
import requests # this is imported within the agent's constructor
monkeypatch.setattr(requests, "post", post)
agent = NomadAgent()
with pytest.raises(ValueError):
agent.deploy_flow(
flow_run=GraphQLResult(
{
"flow": GraphQLResult({"storage": Local().serialize(), "id": "id"}),
"id": "id",
}
)
)
assert not post.called
def test_docker_agent_deploy_flow_uses_environment_metadata(
monkeypatch, cloud_api):
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
api.create_host_config.return_value = {"AutoRemove": True}
monkeypatch.setattr(
"prefect.agent.docker.agent.DockerAgent._get_docker_client",
MagicMock(return_value=api),
)
agent = DockerAgent()
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"id":
"foo",
"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_local_agent_deploy_no_existing_python_path(monkeypatch, cloud_api):
monkeypatch.delenv("PYTHONPATH", raising=False)
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent(import_paths=["paths"])
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local(directory="test").serialize(),
"id": "foo"
}),
"id":
"id",
}))
assert popen.called
assert len(agent.processes) == 1
assert "paths" in popen.call_args[1]["env"]["PYTHONPATH"]
def test_local_agent_deploy_null_or_univeral_run_config(
monkeypatch, run_config):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
agent = LocalAgent()
agent.deploy_flow(flow_run=GraphQLResult(
{
"id": "id",
"flow": {
"storage": Local().serialize(),
"run_config": run_config.serialize() if run_config else None,
"id": "foo",
"core_version": "0.13.0",
},
}, ))
assert popen.called
assert len(agent.processes) == 1
def test_run_flow_calls_callbacks(monkeypatch, tmpdir):
start_func = MagicMock()
exit_func = MagicMock()
file_path = os.path.dirname(
prefect.environments.execution.dask.k8s.__file__)
environment = KubernetesJobEnvironment(os.path.join(file_path, "job.yaml"),
on_start=start_func,
on_exit=exit_func)
flow_runner = MagicMock()
monkeypatch.setattr(
"prefect.engine.get_default_flow_runner_class",
MagicMock(return_value=flow_runner),
)
d = Local(str(tmpdir))
d.add_flow(prefect.Flow("name"))
gql_return = MagicMock(return_value=MagicMock(data=MagicMock(flow_run=[
GraphQLResult({
"flow":
GraphQLResult({
"name": "name",
"storage": d.serialize()
})
})
], )))
client = MagicMock()
client.return_value.graphql = gql_return
monkeypatch.setattr("prefect.environments.execution.base.Client", client)
with set_temporary_config({"cloud.auth_token":
"test"}), prefect.context({"flow_run_id":
"id"}):
environment.run_flow()
assert flow_runner.call_args[1]["flow"].name == "name"
assert start_func.called
assert exit_func.called
def test_docker_agent_deploy_flow_storage_raises(monkeypatch, runner_token):
monkeypatch.setattr("prefect.agent.agent.Client", MagicMock())
api = MagicMock()
api.ping.return_value = True
api.create_container.return_value = {"Id": "container_id"}
monkeypatch.setattr("prefect.agent.docker.agent.docker.APIClient",
MagicMock(return_value=api))
agent = DockerAgent()
with pytest.raises(ValueError):
agent.deploy_flow(flow_run=GraphQLResult(
{
"flow": GraphQLResult({"storage": Local().serialize()}),
"id": "id",
"name": "name",
"version": "version",
}))
assert not api.pull.called
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="Unsupported RunConfig type: Kubernetes"):
agent.deploy_flow(flow_run=GraphQLResult(
{
"id": "id",
"flow": {
"storage": Local().serialize(),
"run_config": KubernetesRun().serialize(),
"id": "foo",
"core_version": "0.13.0",
},
}, ))
assert not popen.called
assert len(agent.processes) == 0
def test_k8s_agent_deploy_flow_raises(monkeypatch, runner_token):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
batch_client = MagicMock()
batch_client.create_namespaced_job.return_value = {}
monkeypatch.setattr("kubernetes.client.BatchV1Api",
MagicMock(retrurn_value=batch_client))
agent = KubernetesAgent()
with pytest.raises(ValueError):
agent.deploy_flow(flow_run=GraphQLResult(
{
"flow": GraphQLResult({
"storage": Local().serialize(),
"id": "id"
}),
"id": "id",
}))
assert not agent.batch_client.create_namespaced_job.called
def test_k8s_agent_deploy_flow_raises(monkeypatch, cloud_api):
k8s_config = MagicMock()
monkeypatch.setattr("kubernetes.config", k8s_config)
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
def test_docker_agent_deploy_flow_unsupported_run_config(api):
agent = DockerAgent()
with pytest.raises(TypeError,
match="Unsupported RunConfig type: LocalRun"):
agent.deploy_flow(flow_run=GraphQLResult({
"flow":
GraphQLResult({
"storage": Local().serialize(),
"run_config": LocalRun().serialize(),
"id": "foo",
"core_version": "0.13.0",
}),
"id":
"id",
"name":
"name",
"version":
"version",
}))
assert not api.pull.called
