def test_run_flow_calls_callbacks(monkeypatch):
start_func = MagicMock()
exit_func = MagicMock()
file_path = os.path.dirname(
prefect.environments.execution.dask.k8s.__file__)
environment = KubernetesJobEnvironment(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),
)
with tempfile.TemporaryDirectory() as directory:
with open(os.path.join(directory, "flow_env.prefect"), "w+") as env:
storage = Local(directory)
flow = prefect.Flow("test", storage=storage)
flow_path = os.path.join(directory, "flow_env.prefect")
with open(flow_path, "wb") as f:
cloudpickle.dump(flow, f)
with set_temporary_config({"cloud.auth_token": "test"}):
with prefect.context(flow_file_path=os.path.join(
directory, "flow_env.prefect")):
environment.run_flow()
assert flow_runner.call_args[1]["flow"].name == "test"
assert start_func.called
assert exit_func.called
def test_run_flow(monkeypatch, tmpdir, job_spec_file):
environment = KubernetesJobEnvironment(job_spec_file=job_spec_file)
flow_runner = MagicMock()
flow_runner_class = MagicMock(return_value=flow_runner)
monkeypatch.setattr(
"prefect.engine.get_default_flow_runner_class",
MagicMock(return_value=flow_runner_class),
)
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_class.call_args[1]["flow"].name == "name"
assert flow_runner.run.call_args[1]["executor"] is environment.executor
def test_run_flow(monkeypatch):
file_path = os.path.dirname(
prefect.environments.execution.dask.k8s.__file__)
environment = KubernetesJobEnvironment(path.join(file_path, "job.yaml"),
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:
with open(os.path.join(directory, "flow_env.prefect"), "w+") as env:
flow = prefect.Flow("test")
flow_path = os.path.join(directory, "flow_env.prefect")
with open(flow_path, "wb") as f:
cloudpickle.dump(flow, f)
with set_temporary_config({"cloud.auth_token": "test"}):
with prefect.context(flow_file_path=os.path.join(
directory, "flow_env.prefect")):
environment.run_flow()
assert flow_runner.call_args[1]["flow"].name == "test"
assert executor.call_args[1] == {"test": "here"}
def test_run_flow_calls_callbacks(monkeypatch):
start_func = MagicMock()
exit_func = MagicMock()
file_path = os.path.dirname(
prefect.environments.execution.dask.k8s.__file__)
environment = KubernetesJobEnvironment(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),
)
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 start_func.called
assert exit_func.called
