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