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