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
