def test_environment_execute_calls_callbacks():
start_func = MagicMock()
exit_func = MagicMock()
with tempfile.TemporaryDirectory() as directory:
@prefect.task
def add_to_dict():
with open(path.join(directory, "output"), "w") as tmp:
tmp.write("success")
with open(path.join(directory, "flow_env.prefect"), "w+") as env:
flow = prefect.Flow("test", tasks=[add_to_dict])
flow_path = path.join(directory, "flow_env.prefect")
with open(flow_path, "wb") as f:
cloudpickle.dump(flow, f)
environment = RemoteEnvironment(on_start=start_func,
on_exit=exit_func)
storage = Local(directory)
storage.add_flow(flow)
environment.execute(flow)
with open(path.join(directory, "output"), "r") as file:
assert file.read() == "success"
assert start_func.called
assert exit_func.called
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_run_flow_no_flow_run_id_in_context(monkeypatch, tmpdir):
environment = Environment()
d = Local(str(tmpdir))
d.add_flow(Flow("name"))
with set_temporary_config({"cloud.auth_token": "test"}):
with pytest.raises(ValueError):
environment.run_flow()
def test_run_flow_no_flow_run_id_in_context(monkeypatch):
environment = Environment()
with tempfile.TemporaryDirectory() as directory:
d = Local(directory)
d.add_flow(Flow("name"))
with set_temporary_config({"cloud.auth_token": "test"}):
with pytest.raises(ValueError):
environment.run_flow()
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 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_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_add_flow_file_to_storage(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)
with pytest.raises(ValueError):
storage.add_flow(f)
storage = Local(stored_as_script=True, path=full_path)
loc = storage.add_flow(f)
assert loc == full_path
assert f.name in storage
def test_environment_execute():
with tempfile.TemporaryDirectory() as directory:
@prefect.task
def add_to_dict():
with open(path.join(directory, "output"), "w") as tmp:
tmp.write("success")
with open(path.join(directory, "flow_env.prefect"), "w+") as env:
flow = prefect.Flow("test", tasks=[add_to_dict])
flow_path = path.join(directory, "flow_env.prefect")
with open(flow_path, "wb") as f:
cloudpickle.dump(flow, f)
environment = RemoteEnvironment()
storage = Local(directory)
storage.add_flow(flow)
environment.execute(flow=flow)
with open(path.join(directory, "output"), "r") as file:
assert file.read() == "success"
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_build_healthchecks():
with tempfile.TemporaryDirectory() as tmpdir:
s = Local(directory=tmpdir)
flow = Flow("TestFlow")
s.add_flow(flow)
assert s.build()
