import math import uuid from collections.abc import Iterable import pytest from celery import Task from celery.canvas import Signature, StampingVisitor, _chain, _chord, chain, chord, group, signature from celery.exceptions import Ignore class LinkingVisitor(StampingVisitor): def on_signature(self, actual_sig: Signature, **headers) -> dict: link_workflow = chain( group(signature("task1"), signature("task2")), signature("task3"), ) link = signature(f"{actual_sig.name}_link") | link_workflow.clone() actual_sig.link(link) link_error = signature(f"{actual_sig.name}_link_error") | link_workflow.clone() actual_sig.link_error(link_error) return super().on_signature(actual_sig, **headers) class CleanupVisitor(StampingVisitor): def clean_stamps(self, actual_sig: Signature) -> None: if "stamped_headers" in actual_sig.options and actual_sig.options["stamped_headers"]: for stamp in actual_sig.options["stamped_headers"]: if stamp in actual_sig.options: actual_sig.options.pop(stamp) def clean_links(self, actual_sig: Signature) -> None: if "link" in actual_sig.options: actual_sig.options.pop("link") if "link_error" in actual_sig.options: actual_sig.options.pop("link_error") def on_signature(self, actual_sig: Signature, **headers) -> dict: self.clean_stamps(actual_sig) self.clean_links(actual_sig) return super().on_signature(actual_sig, **headers) class BooleanStampingVisitor(StampingVisitor): def on_signature(self, actual_sig: Signature, **headers) -> dict: return {"on_signature": True} def on_group_start(self, actual_sig: Signature, **headers) -> dict: return {"on_group_start": True} def on_chain_start(self, actual_sig: Signature, **headers) -> dict: return {"on_chain_start": True} def on_chord_header_start(self, actual_sig: Signature, **header) -> dict: s = super().on_chord_header_start(actual_sig, **header) s.update({"on_chord_header_start": True}) return s def on_chord_body(self, actual_sig: Signature, **header) -> dict: return {"on_chord_body": True} def on_callback(self, actual_sig: Signature, **header) -> dict: return {"on_callback": True} def on_errback(self, actual_sig: Signature, **header) -> dict: return {"on_errback": True} class ListStampingVisitor(StampingVisitor): def on_signature(self, actual_sig: Signature, **headers) -> dict: return { "on_signature": ["ListStampingVisitor: on_signature-item1", "ListStampingVisitor: on_signature-item2"] } def on_group_start(self, actual_sig: Signature, **headers) -> dict: return { "on_group_start": [ "ListStampingVisitor: on_group_start-item1", "ListStampingVisitor: on_group_start-item2", ] } def on_chain_start(self, actual_sig: Signature, **headers) -> dict: return { "on_chain_start": [ "ListStampingVisitor: on_chain_start-item1", "ListStampingVisitor: on_chain_start-item2", ] } def on_chord_header_start(self, actual_sig: Signature, **header) -> dict: s = super().on_chord_header_start(actual_sig, **header) s.update( { "on_chord_header_start": [ "ListStampingVisitor: on_chord_header_start-item1", "ListStampingVisitor: on_chord_header_start-item2", ] } ) return s def on_chord_body(self, actual_sig: Signature, **header) -> dict: return { "on_chord_body": ["ListStampingVisitor: on_chord_body-item1", "ListStampingVisitor: on_chord_body-item2"] } def on_callback(self, actual_sig: Signature, **header) -> dict: return {"on_callback": ["ListStampingVisitor: on_callback-item1", "ListStampingVisitor: on_callback-item2"]} def on_errback(self, actual_sig: Signature, **header) -> dict: return {"on_errback": ["ListStampingVisitor: on_errback-item1", "ListStampingVisitor: on_errback-item2"]} class SetStampingVisitor(StampingVisitor): def on_signature(self, actual_sig: Signature, **headers) -> dict: return { "on_signature": { "SetStampingVisitor: on_signature-item1", "SetStampingVisitor: on_signature-item2", "SetStampingVisitor: on_signature-item3", } } def on_group_start(self, actual_sig: Signature, **headers) -> dict: return { "on_group_start": { "SetStampingVisitor: on_group_start-item1", "SetStampingVisitor: on_group_start-item2", "SetStampingVisitor: on_group_start-item3", } } def on_chain_start(self, actual_sig: Signature, **headers) -> dict: return { "on_chain_start": { "SetStampingVisitor: on_chain_start-item1", "SetStampingVisitor: on_chain_start-item2", "SetStampingVisitor: on_chain_start-item3", } } def on_chord_header_start(self, actual_sig: Signature, **header) -> dict: s = super().on_chord_header_start(actual_sig, **header) s.update( { "on_chord_header_start": { "SetStampingVisitor: on_chord_header_start-item1", "SetStampingVisitor: on_chord_header_start-item2", "SetStampingVisitor: on_chord_header_start-item3", } } ) return s def on_chord_body(self, actual_sig: Signature, **header) -> dict: return { "on_chord_body": { "SetStampingVisitor: on_chord_body-item1", "SetStampingVisitor: on_chord_body-item2", "SetStampingVisitor: on_chord_body-item3", } } def on_callback(self, actual_sig: Signature, **header) -> dict: return { "on_callback": { "SetStampingVisitor: on_callback-item1", "SetStampingVisitor: on_callback-item2", "SetStampingVisitor: on_callback-item3", } } def on_errback(self, actual_sig: Signature, **header) -> dict: return { "on_errback": { "SetStampingVisitor: on_errback-item1", "SetStampingVisitor: on_errback-item2", "SetStampingVisitor: on_errback-item3", } } class StringStampingVisitor(StampingVisitor): def on_signature(self, actual_sig: Signature, **headers) -> dict: return {"on_signature": "StringStampingVisitor: on_signature-item1"} def on_group_start(self, actual_sig: Signature, **headers) -> dict: return {"on_group_start": "StringStampingVisitor: on_group_start-item1"} def on_chain_start(self, actual_sig: Signature, **headers) -> dict: return {"on_chain_start": "StringStampingVisitor: on_chain_start-item1"} def on_chord_header_start(self, actual_sig: Signature, **header) -> dict: s = super().on_chord_header_start(actual_sig, **header) s.update({"on_chord_header_start": "StringStampingVisitor: on_chord_header_start-item1"}) return s def on_chord_body(self, actual_sig: Signature, **header) -> dict: return {"on_chord_body": "StringStampingVisitor: on_chord_body-item1"} def on_callback(self, actual_sig: Signature, **header) -> dict: return {"on_callback": "StringStampingVisitor: on_callback-item1"} def on_errback(self, actual_sig: Signature, **header) -> dict: return {"on_errback": "StringStampingVisitor: on_errback-item1"} class UUIDStampingVisitor(StampingVisitor): frozen_uuid = str(uuid.uuid4()) def on_signature(self, actual_sig: Signature, **headers) -> dict: return {"on_signature": UUIDStampingVisitor.frozen_uuid} def on_group_start(self, actual_sig: Signature, **headers) -> dict: return {"on_group_start": UUIDStampingVisitor.frozen_uuid} def on_chain_start(self, actual_sig: Signature, **headers) -> dict: return {"on_chain_start": UUIDStampingVisitor.frozen_uuid} def on_chord_header_start(self, actual_sig: Signature, **header) -> dict: s = super().on_chord_header_start(actual_sig, **header) s.update({"on_chord_header_start": UUIDStampingVisitor.frozen_uuid}) return s def on_chord_body(self, actual_sig: Signature, **header) -> dict: return {"on_chord_body": UUIDStampingVisitor.frozen_uuid} def on_callback(self, actual_sig: Signature, **header) -> dict: return {"on_callback": UUIDStampingVisitor.frozen_uuid} def on_errback(self, actual_sig: Signature, **header) -> dict: return {"on_errback": UUIDStampingVisitor.frozen_uuid} class StampsAssertionVisitor(StampingVisitor): """ The canvas stamping mechanism traverses the canvas automatically, so we can ride it to traverse the canvas recursively and assert that all signatures have the correct stamp in options """ def __init__(self, visitor: StampingVisitor, subtests): self.visitor = visitor self.subtests = subtests def assertion_check(self, actual_sig: Signature, method: str, **headers) -> None: if any( [ isinstance(actual_sig, group), isinstance(actual_sig, _chain), isinstance(actual_sig, _chord), ] ): return expected_stamp = getattr(self.visitor, method)(actual_sig, **headers)[method] actual_stamp = actual_sig.options[method] with self.subtests.test(f"Check if {actual_sig} has stamp: {expected_stamp}"): if isinstance(self.visitor, ListStampingVisitor) or isinstance(self.visitor, SetStampingVisitor): assertion_check = all([actual in expected_stamp for actual in actual_stamp]) else: assertion_check = actual_stamp == expected_stamp assertion_error = f"{actual_sig} has stamp {actual_stamp} instead of: {expected_stamp}" assert assertion_check, assertion_error def on_signature(self, actual_sig: Signature, **headers) -> dict: self.assertion_check(actual_sig, "on_signature", **headers) return super().on_signature(actual_sig, **headers) def on_group_start(self, actual_sig: Signature, **headers) -> dict: self.assertion_check(actual_sig, "on_group_start", **headers) return super().on_group_start(actual_sig, **headers) def on_chain_start(self, actual_sig: Signature, **headers) -> dict: self.assertion_check(actual_sig, "on_chain_start", **headers) return super().on_chain_start(actual_sig, **headers) def on_chord_header_start(self, actual_sig: Signature, **header) -> dict: self.assertion_check(actual_sig, "on_chord_header_start", **header) if issubclass(type(actual_sig.tasks), Signature): self.assertion_check(actual_sig.tasks, "on_chord_header_start", **header) return super().on_chord_header_start(actual_sig, **header) def on_chord_body(self, actual_sig: chord, **header) -> dict: self.assertion_check(actual_sig.body, "on_chord_body", **header) return super().on_chord_body(actual_sig, **header) def on_callback(self, actual_link_sig: Signature, **header) -> dict: self.assertion_check(actual_link_sig, "on_callback", **header) return super().on_callback(actual_link_sig, **header) def on_errback(self, actual_linkerr_sig: Signature, **header) -> dict: self.assertion_check(actual_linkerr_sig, "on_errback", **header) return super().on_errback(actual_linkerr_sig, **header) class StampedHeadersAssertionVisitor(StampingVisitor): """ The canvas stamping mechanism traverses the canvas automatically, so we can ride it to traverse the canvas recursively and assert that all signatures have the correct stamp in options["stamped_headers"] """ def __init__(self, visitor: StampingVisitor, subtests): self.visitor = visitor self.subtests = subtests def assertion_check(self, actual_sig: Signature, expected_stamped_header: str) -> None: if any( [ isinstance(actual_sig, group), isinstance(actual_sig, _chain), isinstance(actual_sig, _chord), ] ): with self.subtests.test(f'Check if "stamped_headers" is not in {actual_sig.options}'): assertion_check = "stamped_headers" not in actual_sig.options assertion_error = f"{actual_sig} should not have stamped_headers in options" assert assertion_check, assertion_error return actual_stamped_headers = actual_sig.options["stamped_headers"] with self.subtests.test(f'Check if {actual_sig}["stamped_headers"] has: {expected_stamped_header}'): assertion_check = expected_stamped_header in actual_stamped_headers assertion_error = ( f'{actual_sig}["stamped_headers"] {actual_stamped_headers} does ' f"not contain {expected_stamped_header}" ) assert assertion_check, assertion_error def on_signature(self, actual_sig: Signature, **headers) -> dict: self.assertion_check(actual_sig, "on_signature") return super().on_signature(actual_sig, **headers) def on_group_start(self, actual_sig: Signature, **headers) -> dict: self.assertion_check(actual_sig, "on_group_start") return super().on_group_start(actual_sig, **headers) def on_chain_start(self, actual_sig: Signature, **headers) -> dict: self.assertion_check(actual_sig, "on_chain_start") return super().on_chain_start(actual_sig, **headers) def on_chord_header_start(self, actual_sig: Signature, **header) -> dict: self.assertion_check(actual_sig, "on_chord_header_start") if issubclass(type(actual_sig.tasks), Signature): self.assertion_check(actual_sig.tasks, "on_chord_header_start") return super().on_chord_header_start(actual_sig, **header) def on_chord_body(self, actual_sig: chord, **header) -> dict: self.assertion_check(actual_sig.body, "on_chord_body") return super().on_chord_body(actual_sig, **header) def on_callback(self, actual_link_sig: Signature, **header) -> dict: self.assertion_check(actual_link_sig, "on_callback") return super().on_callback(actual_link_sig, **header) def on_errback(self, actual_linkerr_sig: Signature, **header) -> dict: self.assertion_check(actual_linkerr_sig, "on_errback") return super().on_errback(actual_linkerr_sig, **header) def return_True(*args, **kwargs): return True class CanvasCase: def setup_method(self): @self.app.task(shared=False) def identity(x): return x self.identity = identity @self.app.task(shared=False) def fail(*args): args = ("Task expected to fail",) + args raise Exception(*args) self.fail = fail @self.app.task(shared=False) def add(x, y): return x + y self.add = add @self.app.task(shared=False) def mul(x, y): return x * y self.mul = mul @self.app.task(shared=False) def div(x, y): return x / y self.div = div @self.app.task(shared=False) def xsum(numbers): return sum(sum(num) if isinstance(num, Iterable) else num for num in numbers) self.xsum = xsum @self.app.task(shared=False, bind=True) def replaced(self, x, y): return self.replace(add.si(x, y)) self.replaced = replaced @self.app.task(shared=False, bind=True) def replaced_group(self, x, y): return self.replace(group(add.si(x, y), mul.si(x, y))) self.replaced_group = replaced_group @self.app.task(shared=False, bind=True) def replace_with_group(self, x, y): return self.replace(group(add.si(x, y), mul.si(x, y))) self.replace_with_group = replace_with_group @self.app.task(shared=False, bind=True) def replace_with_chain(self, x, y): return self.replace(group(add.si(x, y) | mul.s(y), add.si(x, y))) self.replace_with_chain = replace_with_chain @self.app.task(shared=False) def xprod(numbers): try: return math.prod(numbers) except AttributeError: # TODO: Drop this backport once # we drop support for Python 3.7 import operator from functools import reduce return reduce(operator.mul, numbers) self.xprod = xprod @self.app.task(bind=True, max_retries=3, iterations=0, shared=False) def retry_task(self, arg1, arg2, kwarg=1, max_retries=None, care=True): self.iterations += 1 rmax = self.max_retries if max_retries is None else max_retries assert repr(self.request) retries = self.request.retries if care and retries >= rmax: return arg1 else: raise self.retry(countdown=0, max_retries=rmax) self.retry_task = retry_task @pytest.mark.parametrize( "stamping_visitor", [ BooleanStampingVisitor(), ListStampingVisitor(), SetStampingVisitor(), StringStampingVisitor(), UUIDStampingVisitor(), ], ) @pytest.mark.parametrize( "canvas_workflow", [ signature("sig"), group(signature("sig")), group(signature("sig1", signature("sig2"))), group(signature(f"sig{i}") for i in range(2)), chord((signature(f"sig{i}") for i in range(2)), signature("sig3")), chord(group(signature(f"sig{i}") for i in range(2)), signature("sig3")), chord(group(signature(f"sig{i}") for i in range(2)), signature("sig3") | signature("sig4")), chord(signature("sig1"), signature("sig2") | signature("sig3")), chain( signature("sig"), chord((signature(f"sig{i}") for i in range(2)), signature("sig3")), chord(group(signature(f"sig{i}") for i in range(2)), signature("sig3")), chord(group(signature(f"sig{i}") for i in range(2)), signature("sig3") | signature("sig4")), chord(signature("sig1"), signature("sig2") | signature("sig3")), ), chain( signature("sig1") | signature("sig2"), group(signature("sig3"), signature("sig4")) | group(signature(f"sig{i}") for i in range(5, 6)), chord(group(signature(f"sig{i}") for i in range(6, 8)), signature("sig8")) | signature("sig9"), ), chain( signature("sig"), chord( group(signature(f"sig{i}") for i in range(2)), chain( signature("sig3"), chord( (signature(f"sig{i}") for i in range(4, 6)), chain( signature("sig6"), chord( group(signature(f"sig{i}") for i in range(7, 9)), chain( signature("sig9"), chord(group(signature("sig10"), signature("sig11")), signature("sig12")), ), ), ), ), ), ), ), group( signature("sig"), group(signature("sig1")), group(signature("sig1"), signature("sig2")), group(signature(f"sig{i}") for i in range(2)), group([signature("sig1"), signature("sig2")]), group((signature("sig1"), signature("sig2"))), chain(signature("sig1"), signature("sig2")), chord(group(signature("sig1"), signature("sig2")), signature("sig3")), chord(group(signature(f"sig{i}") for i in range(2)), group(signature("sig3"), signature("sig4"))), chain( group(signature("sig1"), signature("sig2")), group(signature("sig3"), signature("sig4")), signature("sig5"), ), chain( signature("sig1"), group(signature("sig2"), signature("sig3")), group(signature("sig4"), signature("sig5")), ), chain( group( signature("sig1"), group(signature("sig2")), group([signature("sig3"), signature("sig4")]), group(signature(f"sig{i}") for i in range(5, 7)), ), chain( signature("sig8"), group(signature("sig9"), signature("sig10")), ), ), ), chain( signature("sig"), group(signature("sig1")), group(signature("sig1"), signature("sig2")), group(signature(f"sig{i}") for i in range(2)), group([signature("sig1"), signature("sig2")]), group((signature("sig1"), signature("sig2"))), chain(signature("sig1"), signature("sig2")), chord(group(signature("sig1"), signature("sig2")), signature("sig3")), chord(group(signature(f"sig{i}") for i in range(2)), group(signature("sig3"), signature("sig4"))), chain( group(signature("sig1"), signature("sig2")), group(signature("sig3"), signature("sig4")), signature("sig5"), ), chain( signature("sig1"), group(signature("sig2"), signature("sig3")), group(signature("sig4"), signature("sig5")), ), chain( group( signature("sig1"), group(signature("sig2")), group([signature("sig3"), signature("sig4")]), group(signature(f"sig{i}") for i in range(5, 7)), ), chain( signature("sig8"), group(signature("sig9"), signature("sig10")), ), ), ), chord( group( group(signature(f"sig{i}") for i in range(2)), group(signature(f"sig{i}") for i in range(2, 4)), group(signature(f"sig{i}") for i in range(4, 6)), group(signature(f"sig{i}") for i in range(6, 8)), ), chain( chain( signature("sig8") | signature("sig9"), group(signature("sig10"), signature("sig11")) | group(signature(f"sig{i}") for i in range(12, 14)), chord(group(signature(f"sig{i}") for i in range(14, 16)), signature("sig16")) | signature("sig17"), ), signature("sig1") | signature("sig2"), group(signature("sig3"), signature("sig4")) | group(signature(f"sig{i}") for i in range(5, 7)), chord(group(signature(f"sig{i}") for i in range(7, 9)), signature("sig9")) | signature("sig10"), ), ), ], ) class test_canvas_stamping(CanvasCase): @pytest.fixture def stamped_canvas(self, stamping_visitor: StampingVisitor, canvas_workflow: Signature) -> Signature: workflow = canvas_workflow.clone() workflow.stamp(CleanupVisitor()) workflow.stamp(stamping_visitor, append_stamps=False) return workflow @pytest.fixture def stamped_linked_canvas(self, stamping_visitor: StampingVisitor, canvas_workflow: Signature) -> Signature: workflow = canvas_workflow.clone() workflow.stamp(CleanupVisitor()) workflow.stamp(LinkingVisitor()) workflow.stamp(stamping_visitor, append_stamps=False) return workflow @pytest.fixture(params=["stamped_canvas", "stamped_linked_canvas"]) def workflow(self, request, canvas_workflow: Signature) -> Signature: return request.getfixturevalue(request.param) @pytest.mark.usefixtures("depends_on_current_app") def test_stamp_in_options(self, workflow: Signature, stamping_visitor: StampingVisitor, subtests): """Test that all canvas signatures gets the stamp in options""" workflow.stamp(StampsAssertionVisitor(stamping_visitor, subtests)) @pytest.mark.usefixtures("depends_on_current_app") def test_stamping_headers_in_options(self, workflow: Signature, stamping_visitor: StampingVisitor, subtests): """Test that all canvas signatures gets the stamp in options["stamped_headers"]""" workflow.stamp(StampedHeadersAssertionVisitor(stamping_visitor, subtests)) @pytest.mark.usefixtures("depends_on_current_app") def test_stamping_with_replace(self, workflow: Signature, stamping_visitor: StampingVisitor, subtests): class AssertionTask(Task): def on_replace(self, sig: Signature): nonlocal assertion_result assertion_result = True return super().on_replace(sig) @self.app.task(shared=False, bind=True, base=AssertionTask) def assert_using_replace(self: AssertionTask): assert self.request.stamped_headers is None, "stamped_headers should not pass via replace" assert self.request.stamps is None, "stamps should not pass via replace" return self.replace(workflow) @self.app.task(shared=False, bind=True) def stamp_using_replace(self: Task): assert self.request.stamped_headers is not None assert self.request.stamps is not None return self.replace(assert_using_replace.s()) replaced_sig = stamp_using_replace.s() replaced_sig.stamp(stamping_visitor, append_stamps=False) assertion_result = False replaced_sig.apply() assert assertion_result class test_stamping_mechanism(CanvasCase): """These tests were extracted (and fixed) from the canvas unit tests.""" def test_on_signature_gets_the_signature(self): expected_sig = self.add.s(4, 2) class CustomStampingVisitor(StampingVisitor): def on_signature(self, actual_sig, **headers) -> dict: nonlocal expected_sig assert actual_sig == expected_sig return {"header": "value"} sig = expected_sig.clone() sig.stamp(CustomStampingVisitor()) assert sig.options["header"] == "value" def test_double_stamping(self, subtests): """ Test manual signature stamping with two different stamps. """ self.app.conf.task_always_eager = True self.app.conf.task_store_eager_result = True self.app.conf.result_extended = True sig_1 = self.add.s(2, 2) sig_1.stamp(stamp1="stamp1") sig_1.stamp(stamp2="stamp2") sig_1_res = sig_1.freeze() sig_1.apply() with subtests.test("sig_1_res is stamped with stamp1", stamp1=["stamp1"]): assert sig_1_res._get_task_meta()["stamp1"] == ["stamp1"] with subtests.test("sig_1_res is stamped with stamp2", stamp2=["stamp2"]): assert sig_1_res._get_task_meta()["stamp2"] == ["stamp2"] with subtests.test("sig_1_res is stamped twice", stamped_headers=["stamp2", "stamp1"]): assert sorted(sig_1_res._get_task_meta()["stamped_headers"]) == sorted(["stamp2", "stamp1"]) def test_twice_stamping(self, subtests): """ Test manual signature stamping with two stamps twice. """ self.app.conf.task_always_eager = True self.app.conf.task_store_eager_result = True self.app.conf.result_extended = True sig_1 = self.add.s(2, 2) sig_1.stamp(stamp1="stamp1") sig_1.stamp(stamp2="stamp") sig_1.stamp(stamp2="stamp2", append_stamps=True) sig_1.stamp(stamp3=["stamp3"]) sig_1_res = sig_1.freeze() sig_1.apply() with subtests.test("sig_1_res is stamped twice", stamps=["stamp2", "stamp1"]): assert sorted(sig_1_res._get_task_meta()["stamp1"]) == ["stamp1"] assert sorted(sig_1_res._get_task_meta()["stamp2"]) == sorted(["stamp", "stamp2"]) assert sorted(sig_1_res._get_task_meta()["stamp3"]) == ["stamp3"] with subtests.test("sig_1_res is stamped twice", stamped_headers=["stamp2", "stamp1"]): assert sorted(sig_1_res._get_task_meta()["stamped_headers"]) == sorted(["stamp1", "stamp2", "stamp3"]) def test_manual_stamping(self): """ Test manual signature stamping. """ self.app.conf.task_always_eager = True self.app.conf.task_store_eager_result = True self.app.conf.result_extended = True sig_1 = self.add.s(2, 2) stamps = ["stamp1", "stamp2"] sig_1.stamp(visitor=None, groups=[stamps[1]]) sig_1.stamp(visitor=None, groups=stamps[0], append_stamps=True) sig_1_res = sig_1.freeze() sig_1.apply() assert sorted(sig_1_res._get_task_meta()["groups"]) == sorted(stamps) def test_custom_stamping_visitor(self, subtests): """ Test manual signature stamping with a custom visitor class. """ self.app.conf.task_always_eager = True self.app.conf.task_store_eager_result = True self.app.conf.result_extended = True class CustomStampingVisitor1(StampingVisitor): def on_signature(self, sig, **headers) -> dict: # without using stamped_headers key explicitly # the key will be calculated from the headers implicitly return {"header": "value"} class CustomStampingVisitor2(StampingVisitor): def on_signature(self, sig, **headers) -> dict: return {"header": "value", "stamped_headers": ["header"]} sig_1 = self.add.s(2, 2) sig_1.stamp(visitor=CustomStampingVisitor1()) sig_1_res = sig_1.freeze() sig_1.apply() sig_2 = self.add.s(2, 2) sig_2.stamp(visitor=CustomStampingVisitor2()) sig_2_res = sig_2.freeze() sig_2.apply() with subtests.test("sig_1 is stamped with custom visitor", stamped_headers=["header"]): assert sorted(sig_1_res._get_task_meta()["stamped_headers"]) == sorted(["header"]) with subtests.test("sig_2 is stamped with custom visitor", stamped_headers=["header"]): assert sorted(sig_2_res._get_task_meta()["stamped_headers"]) == sorted(["header"]) with subtests.test("sig_1 is stamped with custom visitor", header=["value"]): assert sig_1_res._get_task_meta()["header"] == ["value"] with subtests.test("sig_2 is stamped with custom visitor", header=["value"]): assert sig_2_res._get_task_meta()["header"] == ["value"] def test_callback_stamping(self, subtests): self.app.conf.task_always_eager = True self.app.conf.task_store_eager_result = True self.app.conf.result_extended = True class CustomStampingVisitor(StampingVisitor): def on_signature(self, sig, **headers) -> dict: return {"header": "value"} def on_callback(self, callback, **header) -> dict: return {"on_callback": True} def on_errback(self, errback, **header) -> dict: return {"on_errback": True} sig_1 = self.add.s(0, 1) sig_1_res = sig_1.freeze() group_sig = group([self.add.s(3), self.add.s(4)]) group_sig_res = group_sig.freeze() chord_sig = chord([self.xsum.s(), self.xsum.s()], self.xsum.s()) chord_sig_res = chord_sig.freeze() sig_2 = self.add.s(2) sig_2_res = sig_2.freeze() chain_sig = chain( sig_1, # --> 1 group_sig, # --> [1+3, 1+4] --> [4, 5] chord_sig, # --> [4+5, 4+5] --> [9, 9] --> 9+9 --> 18 sig_2, # --> 18 + 2 --> 20 ) callback = signature("callback_task") errback = signature("errback_task") chain_sig.link(callback) chain_sig.link_error(errback) chain_sig.stamp(visitor=CustomStampingVisitor()) chain_sig_res = chain_sig.apply_async() chain_sig_res.get() with subtests.test("Confirm the chain was executed correctly", result=20): # Before we run our assertions, let's confirm the base functionality of the chain is working # as expected including the links stamping. assert chain_sig_res.result == 20 with subtests.test("sig_1 is stamped with custom visitor", stamped_headers=["header"]): assert sorted(sig_1_res._get_task_meta()["stamped_headers"]) == sorted(["header"]) with subtests.test("group_sig is stamped with custom visitor", stamped_headers=["header"]): for result in group_sig_res.results: assert sorted(result._get_task_meta()["stamped_headers"]) == sorted(["header"]) with subtests.test("chord_sig is stamped with custom visitor", stamped_headers=["header"]): assert sorted(chord_sig_res._get_task_meta()["stamped_headers"]) == sorted(["header"]) with subtests.test("sig_2 is stamped with custom visitor", stamped_headers=["header"]): assert sorted(sig_2_res._get_task_meta()["stamped_headers"]) == sorted(["header"]) with subtests.test( "callback is stamped with custom visitor", stamped_headers=["header", "on_callback"], ): callback_link = chain_sig.options["link"][0] headers = callback_link.options stamped_headers = headers["stamped_headers"] assert sorted(stamped_headers) == sorted(["header", "on_callback"]) assert headers["on_callback"] is True assert headers["header"] == "value" with subtests.test( "errback is stamped with custom visitor", stamped_headers=["header", "on_errback"], ): errback_link = chain_sig.options["link_error"][0] headers = errback_link.options stamped_headers = headers["stamped_headers"] assert sorted(stamped_headers) == sorted(["header", "on_errback"]) assert headers["on_errback"] is True assert headers["header"] == "value" def test_callback_stamping_link_after_stamp(self, subtests): self.app.conf.task_always_eager = True self.app.conf.task_store_eager_result = True self.app.conf.result_extended = True class CustomStampingVisitor(StampingVisitor): def on_signature(self, sig, **headers) -> dict: return {"header": "value"} def on_callback(self, callback, **header) -> dict: return {"on_callback": True} def on_errback(self, errback, **header) -> dict: return {"on_errback": True} sig_1 = self.add.s(0, 1) sig_1_res = sig_1.freeze() group_sig = group([self.add.s(3), self.add.s(4)]) group_sig_res = group_sig.freeze() chord_sig = chord([self.xsum.s(), self.xsum.s()], self.xsum.s()) chord_sig_res = chord_sig.freeze() sig_2 = self.add.s(2) sig_2_res = sig_2.freeze() chain_sig = chain( sig_1, # --> 1 group_sig, # --> [1+3, 1+4] --> [4, 5] chord_sig, # --> [4+5, 4+5] --> [9, 9] --> 9+9 --> 18 sig_2, # --> 18 + 2 --> 20 ) callback = signature("callback_task") errback = signature("errback_task") chain_sig.stamp(visitor=CustomStampingVisitor()) chain_sig.link(callback) chain_sig.link_error(errback) chain_sig_res = chain_sig.apply_async() chain_sig_res.get() with subtests.test("Confirm the chain was executed correctly", result=20): # Before we run our assertions, let's confirm the base functionality of the chain is working # as expected including the links stamping. assert chain_sig_res.result == 20 with subtests.test("sig_1 is stamped with custom visitor", stamped_headers=["header"]): assert sorted(sig_1_res._get_task_meta()["stamped_headers"]) == sorted(["header"]) with subtests.test("group_sig is stamped with custom visitor", stamped_headers=["header"]): for result in group_sig_res.results: assert sorted(result._get_task_meta()["stamped_headers"]) == sorted(["header"]) with subtests.test("chord_sig is stamped with custom visitor", stamped_headers=["header"]): assert sorted(chord_sig_res._get_task_meta()["stamped_headers"]) == sorted(["header"]) with subtests.test("sig_2 is stamped with custom visitor", stamped_headers=["header"]): assert sorted(sig_2_res._get_task_meta()["stamped_headers"]) == sorted(["header"]) with subtests.test("callback is not stamped"): callback_link = chain_sig.options["link"][0] headers = callback_link.options stamped_headers = headers.get("stamped_headers", []) assert "on_callback" not in stamped_headers, "Linking after stamping should not stamp the callback" assert stamped_headers == [] with subtests.test("errback is not stamped"): errback_link = chain_sig.options["link_error"][0] headers = errback_link.options stamped_headers = headers.get("stamped_headers", []) assert "on_callback" not in stamped_headers, "Linking after stamping should not stamp the errback" assert stamped_headers == [] def test_callback_stamping_link_multiple_visitors(self, subtests): self.app.conf.task_always_eager = True self.app.conf.task_store_eager_result = True self.app.conf.result_extended = True class CustomStampingVisitor(StampingVisitor): def on_signature(self, sig, **headers) -> dict: return {"header": "value"} def on_callback(self, callback, **header) -> dict: return {"on_callback": True} def on_errback(self, errback, **header) -> dict: return {"on_errback": True} class CustomStampingVisitor2(StampingVisitor): def on_signature(self, sig, **headers) -> dict: return {"header2": "value2"} def on_callback(self, callback, **header) -> dict: return {"on_callback2": "True"} def on_errback(self, errback, **header) -> dict: return {"on_errback2": "True"} sig_1 = self.add.s(0, 1) sig_1_res = sig_1.freeze() group_sig = group([self.add.s(3), self.add.s(4)]) group_sig_res = group_sig.freeze() chord_sig = chord([self.xsum.s(), self.xsum.s()], self.xsum.s()) chord_sig_res = chord_sig.freeze() sig_2 = self.add.s(2) sig_2_res = sig_2.freeze() chain_sig = chain( sig_1, # --> 1 group_sig, # --> [1+3, 1+4] --> [4, 5] chord_sig, # --> [4+5, 4+5] --> [9, 9] --> 9+9 --> 18 sig_2, # --> 18 + 2 --> 20 ) callback = signature("callback_task") errback = signature("errback_task") chain_sig.stamp(visitor=CustomStampingVisitor()) chain_sig.link(callback) chain_sig.link_error(errback) chain_sig.stamp(visitor=CustomStampingVisitor2()) chain_sig_res = chain_sig.apply_async() chain_sig_res.get() with subtests.test("Confirm the chain was executed correctly", result=20): # Before we run our assertions, let's confirm the base functionality of the chain is working # as expected including the links stamping. assert chain_sig_res.result == 20 with subtests.test("sig_1 is stamped with custom visitor", stamped_headers=["header", "header2"]): assert sorted(sig_1_res._get_task_meta()["stamped_headers"]) == sorted(["header", "header2"]) with subtests.test("group_sig is stamped with custom visitor", stamped_headers=["header", "header2"]): for result in group_sig_res.results: assert sorted(result._get_task_meta()["stamped_headers"]) == sorted(["header", "header2"]) with subtests.test("chord_sig is stamped with custom visitor", stamped_headers=["header", "header2"]): assert sorted(chord_sig_res._get_task_meta()["stamped_headers"]) == sorted(["header", "header2"]) with subtests.test("sig_2 is stamped with custom visitor", stamped_headers=["header", "header2"]): assert sorted(sig_2_res._get_task_meta()["stamped_headers"]) == sorted(["header", "header2"]) with subtests.test("callback is stamped"): callback_link = chain_sig.options["link"][0] headers = callback_link.options stamped_headers = headers.get("stamped_headers", []) assert "on_callback2" in stamped_headers, "Linking after stamping should stamp the callback" expected_stamped_headers = list(CustomStampingVisitor2().on_signature(None).keys()) expected_stamped_headers.extend(list(CustomStampingVisitor2().on_callback(None).keys())) assert sorted(stamped_headers) == sorted(expected_stamped_headers) with subtests.test("errback is stamped"): errback_link = chain_sig.options["link_error"][0] headers = errback_link.options stamped_headers = headers.get("stamped_headers", []) assert "on_errback2" in stamped_headers, "Linking after stamping should stamp the errback" expected_stamped_headers = list(CustomStampingVisitor2().on_signature(None).keys()) expected_stamped_headers.extend(list(CustomStampingVisitor2().on_errback(None).keys())) assert sorted(stamped_headers) == sorted(expected_stamped_headers) @pytest.mark.usefixtures("depends_on_current_app") def test_callback_stamping_on_replace(self, subtests): class CustomStampingVisitor(StampingVisitor): def on_signature(self, sig, **headers) -> dict: return {"header": "value"} def on_callback(self, callback, **header) -> dict: return {"on_callback": True} def on_errback(self, errback, **header) -> dict: return {"on_errback": True} class MyTask(Task): def on_replace(self, sig): sig.stamp(CustomStampingVisitor()) return super().on_replace(sig) mytask = self.app.task(shared=False, base=MyTask)(return_True) sig1 = signature("sig1") callback = signature("callback_task") errback = signature("errback_task") sig1.link(callback) sig1.link_error(errback) with subtests.test("callback is not stamped with custom visitor yet"): callback_link = sig1.options["link"][0] headers = callback_link.options assert "on_callback" not in headers assert "header" not in headers with subtests.test("errback is not stamped with custom visitor yet"): errback_link = sig1.options["link_error"][0] headers = errback_link.options assert "on_errback" not in headers assert "header" not in headers with pytest.raises(Ignore): mytask.replace(sig1) with subtests.test( "callback is stamped with custom visitor", stamped_headers=["header", "on_callback"], ): callback_link = sig1.options["link"][0] headers = callback_link.options stamped_headers = headers["stamped_headers"] assert sorted(stamped_headers) == sorted(["header", "on_callback"]) assert headers["on_callback"] is True assert headers["header"] == "value" with subtests.test( "errback is stamped with custom visitor", stamped_headers=["header", "on_errback"], ): errback_link = sig1.options["link_error"][0] headers = errback_link.options stamped_headers = headers["stamped_headers"] assert sorted(stamped_headers) == sorted(["header", "on_errback"]) assert headers["on_errback"] is True assert headers["header"] == "value" @pytest.mark.parametrize( "sig_to_replace", [ group(signature(f"sig{i}") for i in range(2)), group([signature("sig1"), signature("sig2")]), group((signature("sig1"), signature("sig2"))), group(signature("sig1"), signature("sig2")), chain(signature("sig1"), signature("sig2")), ], ) @pytest.mark.usefixtures("depends_on_current_app") def test_replacing_stamped_canvas_with_tasks(self, subtests, sig_to_replace): class CustomStampingVisitor(StampingVisitor): def on_signature(self, sig, **headers) -> dict: return {"header": "value"} class MyTask(Task): def on_replace(self, sig): nonlocal assertion_result nonlocal failed_task tasks = sig.tasks.tasks if isinstance(sig.tasks, group) else sig.tasks assertion_result = len(tasks) == 2 for task in tasks: assertion_result = all( [ assertion_result, "header" in task.options["stamped_headers"], all([header in task.options for header in task.options["stamped_headers"]]), ] ) if not assertion_result: failed_task = task break return super().on_replace(sig) @self.app.task(shared=False, bind=True, base=MyTask) def replace_from_MyTask(self): # Allows easy assertion for the test without using Mock return self.replace(sig_to_replace) sig = replace_from_MyTask.s() sig.stamp(CustomStampingVisitor()) assertion_result = False failed_task = None sig.apply() assert assertion_result, f"Task {failed_task} was not stamped correctly" @pytest.mark.usefixtures("depends_on_current_app") def test_replacing_stamped_canvas_with_tasks_with_links(self): class CustomStampingVisitor(StampingVisitor): def on_signature(self, sig, **headers) -> dict: return {"header": "value"} class MyTask(Task): def on_replace(self, sig): nonlocal assertion_result nonlocal failed_task nonlocal failed_task_link tasks = sig.tasks.tasks if isinstance(sig.tasks, group) else sig.tasks assertion_result = True for task in tasks: links = task.options["link"] links.extend(task.options["link_error"]) for link in links: assertion_result = all( [ assertion_result, all( [ stamped_header in link["options"] for stamped_header in link["options"]["stamped_headers"] ] ), ] ) else: if not assertion_result: failed_task_link = link break assertion_result = all( [ assertion_result, task.options["stamped_headers"]["header"] == "value", all([header in task.options for header in task.options["stamped_headers"]]), ] ) if not assertion_result: failed_task = task break return super().on_replace(sig) @self.app.task(shared=False, bind=True, base=MyTask) def replace_from_MyTask(self): # Allows easy assertion for the test without using Mock return self.replace(sig_to_replace) s1 = chain(signature("foo11"), signature("foo12")) s1.link(signature("link_foo1")) s1.link_error(signature("link_error_foo1")) s2 = chain(signature("foo21"), signature("foo22")) s2.link(signature("link_foo2")) s2.link_error(signature("link_error_foo2")) sig_to_replace = group([s1, s2]) sig = replace_from_MyTask.s() sig.stamp(CustomStampingVisitor()) assertion_result = False failed_task = None failed_task_link = None sig.apply() err_msg = ( f"Task {failed_task} was not stamped correctly" if failed_task else f"Task link {failed_task_link} was not stamped correctly" if failed_task_link else "Assertion failed" ) assert assertion_result, err_msg def test_group_stamping_one_level(self, subtests): self.app.conf.task_always_eager = True self.app.conf.task_store_eager_result = True self.app.conf.result_extended = True sig_1 = self.add.s(2, 2) sig_2 = self.add.s(4, 4) sig_1_res = sig_1.freeze() sig_2_res = sig_2.freeze() g = group(sig_1, sig_2, app=self.app) g.stamp(stamp="stamp") g.apply() with subtests.test("sig_1_res is stamped manually", stamp=["stamp"]): assert sig_1_res._get_task_meta()["stamp"] == ["stamp"] with subtests.test("sig_2_res is stamped manually", stamp=["stamp"]): assert sig_2_res._get_task_meta()["stamp"] == ["stamp"] with subtests.test("sig_1_res has stamped_headers", stamped_headers=["stamp"]): assert sorted(sig_1_res._get_task_meta()["stamped_headers"]) == sorted(["stamp"]) with subtests.test("sig_2_res has stamped_headers", stamped_headers=["stamp"]): assert sorted(sig_2_res._get_task_meta()["stamped_headers"]) == sorted(["stamp"]) def test_chord_stamping_one_level(self, subtests): """ In the case of group within a chord that is from another canvas element, ensure that chord stamps are added correctly when chord are run in parallel. """ self.app.conf.task_always_eager = True self.app.conf.task_store_eager_result = True self.app.conf.result_extended = True sig_1 = self.add.s(2, 2) sig_2 = self.add.s(4, 4) sig_1_res = sig_1.freeze() sig_2_res = sig_2.freeze() sig_sum = self.xsum.s() g = chord([sig_1, sig_2], sig_sum, app=self.app) g.stamp(stamp="stamp") g.freeze() g.apply() with subtests.test("sig_1_res is stamped manually", stamp=["stamp"]): assert sig_1_res._get_task_meta()["stamp"] == ["stamp"] with subtests.test("sig_2_res is stamped manually", stamp=["stamp"]): assert sig_2_res._get_task_meta()["stamp"] == ["stamp"] with subtests.test("sig_1_res has stamped_headers", stamped_headers=["stamp"]): assert sorted(sig_1_res._get_task_meta()["stamped_headers"]) == sorted(["stamp"]) with subtests.test("sig_2_res has stamped_headers", stamped_headers=["stamp"]): assert sorted(sig_2_res._get_task_meta()["stamped_headers"]) == sorted(["stamp"]) def test_retry_stamping(self): self.retry_task.push_request() self.retry_task.request.stamped_headers = ['stamp'] self.retry_task.request.stamps = {'stamp': 'value'} sig = self.retry_task.signature_from_request() assert sig.options['stamped_headers'] == ['stamp'] assert sig.options['stamp'] == 'value' def test_link_error_does_not_duplicate_stamps(self, subtests): class CustomStampingVisitor(StampingVisitor): def on_group_start(self, group, **headers): return {} def on_chain_start(self, chain, **headers): return {} def on_signature(self, sig, **headers): existing_headers = sig.options.get("headers") or {} existing_stamps = existing_headers.get("stamps") or {} existing_stamp = existing_stamps.get("stamp") existing_stamp = existing_stamp or sig.options.get("stamp") if existing_stamp is None: stamp = str(uuid.uuid4()) return {"stamp": stamp} else: assert False, "stamp already exists" def s(n, fail_flag=False): if not fail_flag: return self.identity.si(str(n)) return self.fail.si(str(n)) def tasks(): tasks = [] for i in range(0, 4): fail_flag = False if i: fail_flag = True sig = s(i, fail_flag) sig.link(s(f"link{str(i)}")) sig.link_error(s(f"link_error{str(i)}")) tasks.append(sig) return tasks with subtests.test("group"): canvas = group(tasks()) canvas.link_error(s("group_link_error")) canvas.stamp(CustomStampingVisitor()) with subtests.test("chord header"): self.app.conf.task_allow_error_cb_on_chord_header = True canvas = chord(tasks(), self.identity.si("body"), app=self.app) canvas.link_error(s("group_link_error")) canvas.stamp(CustomStampingVisitor()) with subtests.test("chord body"): self.app.conf.task_allow_error_cb_on_chord_header = False canvas = chord(tasks(), self.identity.si("body"), app=self.app) canvas.link_error(s("group_link_error")) canvas.stamp(CustomStampingVisitor()) with subtests.test("chain"): canvas = chain(tasks()) canvas.link_error(s("chain_link_error")) canvas.stamp(CustomStampingVisitor())