def test_resource_manager_decorator_init(): manager = resource_manager(name="Test")(ResourceManager) assert manager.name == "Test" assert manager.init_task_kwargs["name"] == "Test" assert manager.setup_task_kwargs["name"] == "Test.setup" assert manager.cleanup_task_kwargs["name"] == "Test.cleanup" assert manager.resource_class == ResourceManager
def test_resource_cannot_be_used_with_multiple_flows(): flow = Flow("test") flow2 = Flow("test2") manager = resource_manager(MyResource) with pytest.raises(ValueError, match="Multiple flows"): with manager(flow=flow): inc(1, flow=flow2)
def test_resource_manager_init_overrides(): init_task_kwargs = {"name": "init_name", "tags": ["init"]} setup_task_kwargs = {"name": "setup_name", "tags": ["setup"]} cleanup_task_kwargs = {"name": "cleanup_name", "tags": ["cleanup"]} manager = resource_manager( MyResource, name="Test", init_task_kwargs=init_task_kwargs, setup_task_kwargs=setup_task_kwargs, cleanup_task_kwargs=cleanup_task_kwargs, ) assert manager.name == "Test" assert manager.resource_class == MyResource assert manager.init_task_kwargs == init_task_kwargs assert manager.setup_task_kwargs == setup_task_kwargs assert manager.cleanup_task_kwargs == { "trigger": resource_cleanup_trigger, "skip_on_upstream_skip": False, **cleanup_task_kwargs, } # Copies used assert manager.init_task_kwargs is not init_task_kwargs assert manager.setup_task_kwargs is not setup_task_kwargs assert manager.cleanup_task_kwargs is not cleanup_task_kwargs
def test_resource_manager_generated_flow_structure(api): manager = resource_manager(MyResource) if api == "functional": with Flow("test") as flow: a = inc(1) context = manager(a) with context as resource: b = add(resource, a) c = inc(b) d = inc(2) e = inc(d) f = inc(3) g = inc(f) else: flow = Flow("test") a = inc(1, flow=flow) context = manager(a, flow=flow) with context as resource: b = add(resource, a, flow=flow) c = inc(b, flow=flow) d = inc(2, flow=flow) e = inc(d, flow=flow) f = inc(3, flow=flow) g = inc(f, flow=flow) # task kwargs successfully forwarded to tasks assert context.init_task.name == "MyResource" assert context.setup_task.name == "MyResource.setup" assert context.cleanup_task.name == "MyResource.cleanup" assert not context.cleanup_task.skip_on_upstream_skip # Reference tasks setup properly assert flow.reference_tasks() == {c, e, g} # Check that: # - Tasks with no upstream dependency in the resource context have # the setup task set as an upstream dependency # - Tasks with no downstream dependency in the resource context have # the cleanup task set as a downstream dependency # - All other tasks only have explicit dependencies assert flow.upstream_tasks(a) == set() assert flow.upstream_tasks(context.init_task) == {a} assert flow.upstream_tasks(context.setup_task) == {context.init_task} assert flow.upstream_tasks(b) == {context.setup_task, a} assert flow.upstream_tasks(c) == {b} assert flow.upstream_tasks(d) == {context.setup_task} assert flow.upstream_tasks(e) == {d} assert flow.upstream_tasks(f) == {context.setup_task} assert flow.upstream_tasks(g) == {f} assert flow.upstream_tasks(context.cleanup_task) == { context.init_task, context.setup_task, c, e, f, }
def test_resource_manager_default_init(): manager = resource_manager(MyResource) assert manager.name == "MyResource" assert manager.resource_class == MyResource assert manager.init_task_kwargs == {"name": "MyResource"} assert manager.setup_task_kwargs == {"name": "MyResource.setup"} assert manager.cleanup_task_kwargs == { "name": "MyResource.cleanup", "trigger": resource_cleanup_trigger, "skip_on_upstream_skip": False, }
def test_resource_manager_sets_and_clears_context(): manager = resource_manager(MyResource) with Flow("test"): m1 = manager(1) m2 = manager(2) assert "resource" not in prefect.context with m1: assert prefect.context["resource"] is m1 with m2: assert prefect.context["resource"] is m2 assert prefect.context["resource"] is m1 assert "resource" not in prefect.context
def test_resource_manager_execution_success(): on_setup = MagicMock(return_value=100) on_cleanup = MagicMock() manager = resource_manager(MyResource) with Flow("test") as flow: context = manager(on_setup=on_setup, on_cleanup=on_cleanup) with context as resource: a = inc(resource) inc(a) state = flow.run() assert on_setup.called assert on_cleanup.call_args == ((100, ), {}) assert state.is_successful() for r in state.result.values(): assert r.is_successful()
def test_resource_cleanup_reference_tasks(): manager = resource_manager(MyResource) with Flow("test") as flow: with manager() as resource: a = inc(resource) b = inc(a) c = inc(2) d = inc(b) assert flow.reference_tasks() == {c, d} with Flow("test") as flow: context = manager() with context as resource: a = inc(resource) b = inc(a) c = inc(2) d = inc(b) e = post_cleanup(upstream_tasks=[context.cleanup_task]) assert flow.reference_tasks() == {c, d, e}
def test_resource_manager_errors_no_flow_in_context(): manager = resource_manager(MyResource) with pytest.raises(ValueError, match="Could not infer an active Flow"): with manager(): pass