def test_environment_execute():
class MyExecutor(LocalDaskExecutor):
submit_called = False
def submit(self, *args, **kwargs):
self.submit_called = True
return super().submit(*args, **kwargs)
global_dict = {}
@prefect.task
def add_to_dict():
global_dict["run"] = True
executor = MyExecutor()
environment = LocalEnvironment(executor=executor)
storage = DummyStorage()
flow = prefect.Flow("test",
tasks=[add_to_dict],
environment=environment,
storage=storage)
storage.add_flow(flow)
environment.execute(flow=flow)
assert global_dict.get("run") is True
assert executor.submit_called
def test_environment_execute_with_env_runner_with_kwargs():
class TestStorage(DummyStorage):
def get_flow(self, *args, **kwargs):
raise NotImplementedError()
def get_env_runner(self, flow_loc):
runner = super().get_flow(flow_loc)
def runner_func(env):
runner.run(x=env.get("x"))
return runner_func
global_dict = {}
@prefect.task
def add_to_dict(x):
global_dict["result"] = x
environment = LocalEnvironment()
storage = TestStorage()
with prefect.Flow("test") as flow:
x = prefect.Parameter("x")
add_to_dict(x)
flow_loc = storage.add_flow(flow)
environment.execute(storage, flow_loc, env=dict(x=42))
assert global_dict.get("result") == 42
def test_environment_execute():
global_dict = {}
@prefect.task
def add_to_dict():
global_dict["run"] = True
environment = LocalEnvironment()
storage = DummyStorage()
flow = prefect.Flow("test", tasks=[add_to_dict])
flow_loc = storage.add_flow(flow)
environment.execute(storage, flow_loc)
assert global_dict.get("run") is True
def test_environment_execute_with_kwargs():
global_dict = {}
@prefect.task
def add_to_dict(x):
global_dict["result"] = x
environment = LocalEnvironment()
storage = Memory()
with prefect.Flow("test") as flow:
x = prefect.Parameter("x")
add_to_dict(x)
flow_loc = storage.add_flow(flow)
environment.execute(storage, flow_loc, x=42)
assert global_dict.get("result") == 42
def test_environment_execute_calls_callbacks():
start_func = MagicMock()
exit_func = MagicMock()
global_dict = {}
@prefect.task
def add_to_dict():
global_dict["run"] = True
environment = LocalEnvironment(on_start=start_func, on_exit=exit_func)
storage = DummyStorage()
flow = prefect.Flow("test", tasks=[add_to_dict])
flow_loc = storage.add_flow(flow)
environment.execute(storage, flow_loc)
assert global_dict.get("run") is True
assert start_func.called
assert exit_func.called
def test_environment_execute_with_env_runner():
class TestStorage(DummyStorage):
def get_flow(self, *args, **kwargs):
raise NotImplementedError()
def get_env_runner(self, flow_loc):
runner = super().get_flow(flow_loc)
return lambda env: runner.run()
global_dict = {}
@prefect.task
def add_to_dict():
global_dict["run"] = True
environment = LocalEnvironment()
storage = TestStorage()
flow = prefect.Flow("test", tasks=[add_to_dict])
flow_loc = storage.add_flow(flow)
environment.execute(storage, flow_loc)
assert global_dict.get("run") is True