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():
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():
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_create_environment():
environment = LocalEnvironment()
assert environment
assert environment.labels == set()
assert environment.on_start is None
assert environment.on_exit is None
assert environment.logger.name == "prefect.LocalEnvironment"
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_client_register_flow_id_output(
patch_post, use_run_config, compressed, monkeypatch, capsys, cloud_api, tmpdir
):
if compressed:
response = {
"data": {
"project": [{"id": "proj-id"}],
"create_flow_from_compressed_string": {"id": "long-id"},
"flow_by_pk": {"flow_group_id": "fg-id"},
}
}
else:
response = {
"data": {
"project": [{"id": "proj-id"}],
"create_flow": {"id": "long-id"},
"flow_by_pk": {"flow_group_id": "fg-id"},
}
}
patch_post(response)
monkeypatch.setattr(
"prefect.client.Client.get_default_tenant_slug", MagicMock(return_value="tslug")
)
with set_temporary_config(
{
"cloud.api": "http://my-cloud.foo",
"cloud.auth_token": "secret_token",
"backend": "cloud",
}
):
client = Client()
labels = ["test1", "test2"]
storage = Local(tmpdir)
if use_run_config:
flow = prefect.Flow(
name="test", storage=storage, run_config=LocalRun(labels=labels)
)
flow.environment = None
else:
flow = prefect.Flow(
name="test", storage=storage, environment=LocalEnvironment(labels=labels)
)
flow.result = flow.storage.result
flow_id = client.register(
flow,
project_name="my-default-project",
compressed=compressed,
version_group_id=str(uuid.uuid4()),
)
assert flow_id == "long-id"
captured = capsys.readouterr()
assert "Flow URL: https://cloud.prefect.io/tslug/flow/fg-id\n" in captured.out
assert f"Labels: {labels}" in captured.out
def test_client_register_default_prefect_image(
patch_post, compressed, monkeypatch, tmpdir
):
if compressed:
response = {
"data": {
"project": [{"id": "proj-id"}],
"create_flow_from_compressed_string": {"id": "long-id"},
}
}
else:
response = {
"data": {"project": [{"id": "proj-id"}], "create_flow": {"id": "long-id"}}
}
post = patch_post(response)
monkeypatch.setattr(
"prefect.client.Client.get_default_tenant_slug", MagicMock(return_value="tslug")
)
monkeypatch.setattr("prefect.storage.Docker._build_image", MagicMock())
with set_temporary_config(
{
"cloud.api": "http://my-cloud.foo",
"cloud.auth_token": "secret_token",
"backend": "cloud",
}
):
client = Client()
flow = prefect.Flow(
name="test",
storage=prefect.storage.Local(tmpdir),
environment=LocalEnvironment(),
)
flow.result = flow.storage.result
client.register(
flow,
project_name="my-default-project",
compressed=compressed,
build=True,
no_url=True,
)
# extract POST info
if compressed:
serialized_flow = decompress(
json.loads(post.call_args[1]["json"]["variables"])["input"][
"serialized_flow"
]
)
else:
serialized_flow = json.loads(post.call_args[1]["json"]["variables"])["input"][
"serialized_flow"
]
assert serialized_flow["storage"] is not None
assert "prefecthq/prefect" in serialized_flow["environment"]["metadata"]["image"]
def test_create_environment_populated():
def f():
pass
environment = LocalEnvironment(labels=["test"], on_start=f, on_exit=f)
assert environment
assert environment.labels == set(["test"])
assert environment.on_start is f
assert environment.on_exit is f
assert environment.logger.name == "prefect.LocalEnvironment"
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_create_environment():
with set_temporary_config(
{"engine.executor.default_class": "prefect.executors.LocalDaskExecutor"}
):
environment = LocalEnvironment()
assert isinstance(environment.executor, LocalDaskExecutor)
assert environment.labels == set()
assert environment.on_start is None
assert environment.on_exit is None
assert environment.metadata == {}
assert environment.logger.name == "prefect.LocalEnvironment"
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
def test_create_environment_populated():
def f():
pass
executor = LocalDaskExecutor()
environment = LocalEnvironment(
executor=executor,
labels=["test"],
on_start=f,
on_exit=f,
metadata={"test": "here"},
)
assert environment.executor is executor
assert environment.labels == set(["test"])
assert environment.on_start is f
assert environment.on_exit is f
assert environment.metadata == {"test": "here"}
assert environment.logger.name == "prefect.LocalEnvironment"
def test_serialize_environment():
environment = LocalEnvironment()
env = environment.serialize()
assert env["type"] == "LocalEnvironment"
def test_setup_environment_passes():
environment = LocalEnvironment()
environment.setup(flow=Flow("test", storage=Docker()))
def test_environment_dependencies():
environment = LocalEnvironment()
assert environment.dependencies == []
def test_build_and_register(self, capsys, monkeypatch, force):
"""Build and register a few flows:
- 1 new flow
- 1 updated flow
- 1 skipped flow
- 1 error during registration
- 2 sharing the same storage (which fails to build properly)
"""
build_call_count = 0
class MyModule(Module):
def build(self):
nonlocal build_call_count
build_call_count += 1
class BadStorage(Module):
def build(self):
raise ValueError("whoops!")
client = MagicMock()
client.graphql.side_effect = [
GraphQLResult({"data": {"flow": []}}),
GraphQLResult({"data": {"flow": [{"id": "old-id-2", "version": 1}]}}),
GraphQLResult({"data": {"flow": [{"id": "old-id-3", "version": 2}]}}),
GraphQLResult({"data": {"flow": [{"id": "old-id-4", "version": 3}]}}),
]
client.register.side_effect = [
"new-id-1",
"old-id-2",
"new-id-3",
ValueError("Oh no!"),
]
storage1 = MyModule("testing")
storage1.result = LocalResult()
flow1 = Flow("flow 1", storage=storage1, run_config=UniversalRun(labels=["a"]))
flow2 = Flow(
"flow 2",
storage=MyModule("testing"),
environment=LocalEnvironment(labels=["a"]),
)
storage2 = MyModule("testing")
flow3 = Flow("flow 3", storage=storage2)
flow4 = Flow("flow 4", storage=storage2)
storage3 = BadStorage("testing")
flow5 = Flow("flow 5", storage=storage3)
flow6 = Flow("flow 6", storage=storage3)
flows = [flow1, flow2, flow3, flow4, flow5, flow6]
stats = build_and_register(
client, flows, "testing", labels=["b", "c"], force=force
)
# 3 calls (one for each unique `MyModule` storage object)
assert build_call_count == 3
# 4 register calls (6 - 2 that failed to build storage)
assert client.register.call_count == 4
for flow, (args, kwargs) in zip(flows, client.register.call_args_list):
assert not args
assert kwargs["flow"] is flow
assert kwargs["project_name"] == "testing"
assert kwargs["build"] is False
assert kwargs["no_url"] is True
if force:
assert kwargs["idempotency_key"] is None
else:
assert kwargs["idempotency_key"]
# Stats are recorded properly
assert dict(stats) == {"registered": 2, "skipped": 1, "errored": 3}
# Flows are properly configured
assert flow1.result is storage1.result
assert flow1.run_config.labels == {"a", "b", "c"}
assert flow2.environment.labels == {"a", "b", "c"}
assert isinstance(flow3.run_config, UniversalRun)
assert flow3.run_config.labels == {"b", "c"}
assert isinstance(flow4.run_config, UniversalRun)
assert flow4.run_config.labels == {"b", "c"}
# The output contains a traceback, which will vary between machines
# We only check that the following fixed sections exist in the output
parts = [
(
" Building `MyModule` storage...\n"
" Registering 'flow 1'... Done\n"
" └── ID: new-id-1\n"
" └── Version: 1\n"
" Building `MyModule` storage...\n"
" Registering 'flow 2'... Skipped\n"
" Building `MyModule` storage...\n"
" Registering 'flow 3'... Done\n"
" └── ID: new-id-3\n"
" └── Version: 3\n"
" Registering 'flow 4'... Error\n"
" Traceback (most recent call last):\n"
),
(
" ValueError: Oh no!\n"
"\n"
" Building `BadStorage` storage...\n"
" Error building storage:\n"
" Traceback (most recent call last):\n"
),
(
" ValueError: whoops!\n"
"\n"
" Registering 'flow 5'... Error\n"
" Registering 'flow 6'... Error\n"
),
]
out, err = capsys.readouterr()
assert not err
for part in parts:
assert part in out
def test_build_and_register(self, capsys, monkeypatch, force):
"""Build and register a few flows:
- 1 new flow
- 1 updated flow
- 1 skipped flow
- 1 error during registration
- 2 sharing the same storage (which fails to build properly)
- 2 from a pre-built JSON file
"""
build_call_count = 0
class MyModule(Module):
def build(self):
nonlocal build_call_count
build_call_count += 1
class BadStorage(Module):
def build(self):
raise ValueError("whoops!")
client = MagicMock()
register_serialized_flow = MagicMock()
register_serialized_flow.side_effect = [
("new-id-1", 1, True),
("old-id-2", 2, False),
("new-id-3", 3, True),
ValueError("Oh no!"),
("new-id-7", 1, True),
("old-id-8", 2, False),
]
monkeypatch.setattr(
"prefect.cli.build_register.register_serialized_flow",
register_serialized_flow,
)
storage1 = MyModule("testing")
storage1.result = LocalResult()
flow1 = Flow("flow 1",
storage=storage1,
run_config=UniversalRun(labels=["a"]))
flow2 = Flow(
"flow 2",
storage=MyModule("testing"),
environment=LocalEnvironment(labels=["a"]),
)
storage2 = MyModule("testing")
flow3 = Flow("flow 3", storage=storage2)
flow4 = Flow("flow 4", storage=storage2)
storage3 = BadStorage("testing")
flow5 = Flow("flow 5", storage=storage3)
flow6 = Flow("flow 6", storage=storage3)
flow7 = box.Box(
Flow("flow 7",
run_config=UniversalRun(labels=["a"])).serialize(build=False))
flow8 = box.Box(
Flow("flow 8", environment=LocalEnvironment(
labels=["a"])).serialize(build=False))
flows = [flow1, flow2, flow3, flow4, flow5, flow6, flow7, flow8]
stats = build_and_register(client,
flows,
"my-project-id",
labels=["b", "c"],
force=force)
# 3 calls (one for each unique `MyModule` storage object)
assert build_call_count == 3
# 6 register calls (8 - 2 that failed to build storage)
assert register_serialized_flow.call_count == 6
for flow, (args,
kwargs) in zip(flows,
register_serialized_flow.call_args_list):
assert not args
assert kwargs["client"] is client
assert kwargs["serialized_flow"]
assert kwargs["project_id"] == "my-project-id"
assert kwargs["force"] == force
# Stats are recorded properly
assert dict(stats) == {"registered": 3, "skipped": 2, "errored": 3}
# Flows are properly configured
assert flow1.result is storage1.result
assert flow1.run_config.labels == {"a", "b", "c"}
assert flow2.environment.labels == {"a", "b", "c"}
assert isinstance(flow3.run_config, UniversalRun)
assert flow3.run_config.labels == {"b", "c"}
assert isinstance(flow4.run_config, UniversalRun)
assert flow4.run_config.labels == {"b", "c"}
assert set(flow7["run_config"]["labels"]) == {"a", "b", "c"}
assert set(flow8["environment"]["labels"]) == {"a", "b", "c"}
# The output contains a traceback, which will vary between machines
# We only check that the following fixed sections exist in the output
parts = [
(" Building `MyModule` storage...\n"
" Registering 'flow 1'... Done\n"
" └── ID: new-id-1\n"
" └── Version: 1\n"
" Building `MyModule` storage...\n"
" Registering 'flow 2'... Skipped (metadata unchanged)\n"
" Building `MyModule` storage...\n"
" Registering 'flow 3'... Done\n"
" └── ID: new-id-3\n"
" └── Version: 3\n"
" Registering 'flow 4'... Error\n"
" Traceback (most recent call last):\n"),
(" ValueError: Oh no!\n"
"\n"
" Building `BadStorage` storage...\n"
" Error building storage:\n"
" Traceback (most recent call last):\n"),
(" ValueError: whoops!\n"
"\n"
" Registering 'flow 5'... Error\n"
" Registering 'flow 6'... Error\n"
" Registering 'flow 7'... Done\n"
" └── ID: new-id-7\n"
" └── Version: 1\n"
" Registering 'flow 8'... Skipped (metadata unchanged)\n"),
]
out, err = capsys.readouterr()
assert not err
for part in parts:
assert part in out
def test_create_environment():
environment = LocalEnvironment()
assert environment
def test_setup_environment_passes():
environment = LocalEnvironment()
environment.setup(storage=Docker())
assert environment