From 73eef627b8706ea0bb36ccaf49c452c136685be2 Mon Sep 17 00:00:00 2001 From: Simon Kelly Date: Wed, 22 Jul 2026 11:17:58 +0200 Subject: [PATCH 1/2] Add external_id to Task API and SDK Co-Authored-By: Claude Opus 4.8 (1M context) --- taskbadger.yaml | 15 +++++++++++ .../internal/models/patched_task_request.py | 10 ++++++++ taskbadger/internal/models/task.py | 10 ++++++++ taskbadger/internal/models/task_request.py | 10 ++++++++ taskbadger/sdk.py | 12 +++++++++ tests/test_sdk.py | 25 +++++++++++++++++++ 6 files changed, 82 insertions(+) diff --git a/taskbadger.yaml b/taskbadger.yaml index c9a56e7..750a24b 100644 --- a/taskbadger.yaml +++ b/taskbadger.yaml @@ -694,6 +694,11 @@ components: type: string description: Queue the task is from maxLength: 255 + external_id: + type: string + description: Identifier from the originating system (e.g. Celery task ID) + for correlating with logs + maxLength: 255 status: allOf: - $ref: '#/components/schemas/StatusEnum' @@ -794,6 +799,11 @@ components: type: string description: Queue the task is from maxLength: 255 + external_id: + type: string + description: Identifier from the originating system (e.g. Celery task ID) + for correlating with logs + maxLength: 255 status: allOf: - $ref: '#/components/schemas/StatusEnum' @@ -893,6 +903,11 @@ components: type: string description: Queue the task is from maxLength: 255 + external_id: + type: string + description: Identifier from the originating system (e.g. Celery task ID) + for correlating with logs + maxLength: 255 status: allOf: - $ref: '#/components/schemas/StatusEnum' diff --git a/taskbadger/internal/models/patched_task_request.py b/taskbadger/internal/models/patched_task_request.py index e455f70..40a7d72 100644 --- a/taskbadger/internal/models/patched_task_request.py +++ b/taskbadger/internal/models/patched_task_request.py @@ -24,6 +24,8 @@ class PatchedTaskRequest: Attributes: name (str | Unset): Name of the task queue (str | Unset): Queue the task is from + external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with + logs status (StatusEnum | Unset): * `pending` - pending * `pre_processing` - pre_processing * `processing` - processing @@ -50,6 +52,7 @@ class PatchedTaskRequest: name: str | Unset = UNSET queue: str | Unset = UNSET + external_id: str | Unset = UNSET status: StatusEnum | Unset = StatusEnum.PENDING value: int | None | Unset = UNSET value_max: int | Unset = UNSET @@ -69,6 +72,8 @@ def to_dict(self) -> dict[str, Any]: queue = self.queue + external_id = self.external_id + status: str | Unset = UNSET if not isinstance(self.status, Unset): status = self.status.value @@ -128,6 +133,8 @@ def to_dict(self) -> dict[str, Any]: field_dict["name"] = name if queue is not UNSET: field_dict["queue"] = queue + if external_id is not UNSET: + field_dict["external_id"] = external_id if status is not UNSET: field_dict["status"] = status if value is not UNSET: @@ -160,6 +167,8 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T: queue = d.pop("queue", UNSET) + external_id = d.pop("external_id", UNSET) + _status = d.pop("status", UNSET) status: StatusEnum | Unset if isinstance(_status, Unset): @@ -251,6 +260,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: patched_task_request = cls( name=name, queue=queue, + external_id=external_id, status=status, value=value, value_max=value_max, diff --git a/taskbadger/internal/models/task.py b/taskbadger/internal/models/task.py index 632f325..718616b 100644 --- a/taskbadger/internal/models/task.py +++ b/taskbadger/internal/models/task.py @@ -32,6 +32,8 @@ class Task: url (str): public_url (str): queue (str | Unset): Queue the task is from + external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with + logs status (StatusEnum | Unset): * `pending` - pending * `pre_processing` - pre_processing * `processing` - processing @@ -66,6 +68,7 @@ class Task: url: str public_url: str queue: str | Unset = UNSET + external_id: str | Unset = UNSET status: StatusEnum | Unset = StatusEnum.PENDING value: int | None | Unset = UNSET value_max: int | Unset = UNSET @@ -102,6 +105,8 @@ def to_dict(self) -> dict[str, Any]: queue = self.queue + external_id = self.external_id + status: str | Unset = UNSET if not isinstance(self.status, Unset): status = self.status.value @@ -171,6 +176,8 @@ def to_dict(self) -> dict[str, Any]: ) if queue is not UNSET: field_dict["queue"] = queue + if external_id is not UNSET: + field_dict["external_id"] = external_id if status is not UNSET: field_dict["status"] = status if value is not UNSET: @@ -224,6 +231,8 @@ def _parse_value_percent(data: object) -> int | None: queue = d.pop("queue", UNSET) + external_id = d.pop("external_id", UNSET) + _status = d.pop("status", UNSET) status: StatusEnum | Unset if isinstance(_status, Unset): @@ -323,6 +332,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: url=url, public_url=public_url, queue=queue, + external_id=external_id, status=status, value=value, value_max=value_max, diff --git a/taskbadger/internal/models/task_request.py b/taskbadger/internal/models/task_request.py index 35198ab..363630c 100644 --- a/taskbadger/internal/models/task_request.py +++ b/taskbadger/internal/models/task_request.py @@ -24,6 +24,8 @@ class TaskRequest: Attributes: name (str): Name of the task queue (str | Unset): Queue the task is from + external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with + logs status (StatusEnum | Unset): * `pending` - pending * `pre_processing` - pre_processing * `processing` - processing @@ -50,6 +52,7 @@ class TaskRequest: name: str queue: str | Unset = UNSET + external_id: str | Unset = UNSET status: StatusEnum | Unset = StatusEnum.PENDING value: int | None | Unset = UNSET value_max: int | Unset = UNSET @@ -69,6 +72,8 @@ def to_dict(self) -> dict[str, Any]: queue = self.queue + external_id = self.external_id + status: str | Unset = UNSET if not isinstance(self.status, Unset): status = self.status.value @@ -130,6 +135,8 @@ def to_dict(self) -> dict[str, Any]: ) if queue is not UNSET: field_dict["queue"] = queue + if external_id is not UNSET: + field_dict["external_id"] = external_id if status is not UNSET: field_dict["status"] = status if value is not UNSET: @@ -162,6 +169,8 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T: queue = d.pop("queue", UNSET) + external_id = d.pop("external_id", UNSET) + _status = d.pop("status", UNSET) status: StatusEnum | Unset if isinstance(_status, Unset): @@ -253,6 +262,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: task_request = cls( name=name, queue=queue, + external_id=external_id, status=status, value=value, value_max=value_max, diff --git a/taskbadger/sdk.py b/taskbadger/sdk.py index 4d7fd85..4389cf2 100644 --- a/taskbadger/sdk.py +++ b/taskbadger/sdk.py @@ -152,6 +152,7 @@ def create_task( monitor_id: str = None, tags: dict[str, str] = None, queue: str = None, + external_id: str = None, ) -> "Task": """Create a Task. @@ -167,6 +168,7 @@ def create_task( monitor_id: ID of the monitor to associate this task with. tags: Dictionary of namespace -> value tags. queue: Name of the queue the task is from. + external_id: Identifier from the originating system (e.g. Celery task ID) for correlating with logs. Returns: Task: The created Task object. @@ -177,6 +179,8 @@ def create_task( } if queue is not None: task_dict["queue"] = queue + if external_id is not None: + task_dict["external_id"] = external_id if value is not None: task_dict["value"] = value if value_max is not None: @@ -222,6 +226,7 @@ def update_task( actions: list[Action] = None, tags: dict[str, str] = None, queue: str = None, + external_id: str = None, ) -> "Task": """Update a task. Requires only the task ID and fields to update. @@ -238,6 +243,7 @@ def update_task( actions: Task actions. **Deprecated:** use project-level actions instead. tags: Dictionary of namespace -> value tags. queue: Name of the queue the task is from. + external_id: Identifier from the originating system (e.g. Celery task ID) for correlating with logs. Returns: Task: The updated Task object. @@ -250,6 +256,7 @@ def update_task( max_runtime = _none_to_unset(max_runtime) stale_timeout = _none_to_unset(stale_timeout) queue = _none_to_unset(queue) + external_id = _none_to_unset(external_id) data = data or UNSET body = PatchedTaskRequest( @@ -261,6 +268,7 @@ def update_task( max_runtime=max_runtime, stale_timeout=stale_timeout, queue=queue, + external_id=external_id, ) if actions: _warn_actions_deprecated() @@ -335,6 +343,7 @@ def create( monitor_id: str = None, tags: dict[str, str] = None, queue: str = None, + external_id: str = None, ) -> "Task": """Create a new task @@ -352,6 +361,7 @@ def create( monitor_id=monitor_id, tags=tags, queue=queue, + external_id=external_id, ) def __init__(self, task): @@ -437,6 +447,7 @@ def update( actions: list[Action] = None, tags: dict[str, str] = None, queue: str = None, + external_id: str = None, data_merge_strategy: Any = None, ): """Generic update method used to update any of the task fields. @@ -465,6 +476,7 @@ def update( actions=actions, tags=tags, queue=queue, + external_id=external_id, ) self._task = task._task diff --git a/tests/test_sdk.py b/tests/test_sdk.py index 17878de..768cade 100644 --- a/tests/test_sdk.py +++ b/tests/test_sdk.py @@ -124,6 +124,31 @@ def test_update_queue(settings, patched_update): _verify_update(settings, patched_update, queue="high_priority") +def test_create_with_external_id(settings, patched_create): + api_task = task_for_test() + patched_create.return_value = Response(HTTPStatus.OK, b"", {}, api_task) + + Task.create(name="task name", external_id="celery-abc-123") + + request = TaskRequest(name="task name", status=StatusEnum.PENDING, external_id="celery-abc-123") + patched_create.assert_called_with( + client=mock.ANY, + organization_slug="org", + project_slug="project", + body=request, + ) + + +def test_update_external_id(settings, patched_update): + api_task = task_for_test() + task = Task(api_task) + + patched_update.return_value = Response(HTTPStatus.OK, b"", {}, api_task) + task.update(external_id="celery-abc-123") + + _verify_update(settings, patched_update, external_id="celery-abc-123") + + def test_before_create_update_task(settings, patched_create): def before_create(task): tags = task.setdefault("tags", {}) From 4a4fc3f27408f36ebb25c464ea83e0ae35e9fc35 Mon Sep 17 00:00:00 2001 From: Simon Kelly Date: Wed, 22 Jul 2026 11:18:05 +0200 Subject: [PATCH 2/2] Record celery/procrastinate task ID as external_id Co-Authored-By: Claude Opus 4.8 (1M context) --- taskbadger/celery.py | 8 ++- taskbadger/procrastinate.py | 27 +++++++- tests/test_celery.py | 62 ++++++++++++++++- tests/test_celery_system_integration.py | 7 +- tests/test_procrastinate.py | 92 ++++++++++++++++++++++--- 5 files changed, 178 insertions(+), 18 deletions(-) diff --git a/taskbadger/celery.py b/taskbadger/celery.py index b0508d3..e61964b 100644 --- a/taskbadger/celery.py +++ b/taskbadger/celery.py @@ -137,7 +137,9 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs): ctask = celery.current_app.tasks.get(sender) # get kwargs from the task class (set via decorator) - kwargs = getattr(ctask, TB_KWARGS_ARG, {}) + # copy: the raw attr is shared across invocations, and we mutate per-publish + # values (external_id, status, name) into it below + kwargs = dict(getattr(ctask, TB_KWARGS_ARG, {})) for attr in dir(ctask): if attr.startswith(KWARG_PREFIX) and attr not in IGNORE_ARGS: kwargs[attr.removeprefix(KWARG_PREFIX)] = getattr(ctask, attr) @@ -147,6 +149,7 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs): kwargs["status"] = StatusEnum.PENDING if routing_key and "queue" not in kwargs: kwargs["queue"] = routing_key + kwargs.setdefault("external_id", headers["id"]) name = kwargs.pop("name", headers["task"]) global_record_task_args = celery_system and celery_system.record_task_args @@ -247,7 +250,8 @@ def _maybe_create_task(signal_sender): delivery_info = getattr(signal_sender.request, "delivery_info", None) or {} queue = delivery_info.get("routing_key") - task = create_task_safe(task_name, status=StatusEnum.PENDING, data=data, queue=queue) + external_id = signal_sender.request.id + task = create_task_safe(task_name, status=StatusEnum.PENDING, data=data, queue=queue, external_id=external_id) if task: # Store the task ID in the request so _update_task can find it signal_sender.request.update({TB_TASK_ID: task.id}) diff --git a/taskbadger/procrastinate.py b/taskbadger/procrastinate.py index 7448e27..3ae6364 100644 --- a/taskbadger/procrastinate.py +++ b/taskbadger/procrastinate.py @@ -149,12 +149,16 @@ def _wrap_defer(task): @functools.wraps(original_defer) def defer(**kwargs): kwargs = _maybe_create_pending(task, kwargs) - return original_defer(**kwargs) + job_id = original_defer(**kwargs) + _record_external_id(kwargs, job_id) + return job_id @functools.wraps(original_defer_async) async def defer_async(**kwargs): kwargs = _maybe_create_pending(task, kwargs) - return await original_defer_async(**kwargs) + job_id = await original_defer_async(**kwargs) + _record_external_id(kwargs, job_id) + return job_id task.defer = defer task.defer_async = defer_async @@ -214,6 +218,18 @@ def _maybe_create_pending(task, kwargs): return new_kwargs +def _record_external_id(kwargs, job_id): + """Record the Procrastinate job id as the TaskBadger task's ``external_id``. + + ``defer`` only returns the DB-assigned job id once the job is enqueued, so this + runs as a follow-up update after both ids are known. No-op if the defer wasn't + tracked (no injected id) or no job id came back.""" + tb_id = kwargs.get(TB_TASK_ID_KWARG) + if tb_id is None or job_id is None: + return + update_task_safe(tb_id, external_id=str(job_id)) + + def _serialize_kwargs(kwargs): """Return a JSON-roundtrippable copy of the defer kwargs. @@ -296,14 +312,19 @@ def _patch_job_manager(app, system): @functools.wraps(original) async def patched(*, job, periodic_id, defer_timestamp): task = app.tasks.get(job.task_name) + tb_id = None if task is not None: tb_task = _create_pending_task(task, job.task_kwargs, queue=job.queue) if tb_task is not None: new_kwargs = {**job.task_kwargs, TB_TASK_ID_KWARG: tb_task.id} job = job.evolve(task_kwargs=new_kwargs) - return await jm._taskbadger_original_defer_periodic_job( + tb_id = tb_task.id + job_id = await jm._taskbadger_original_defer_periodic_job( job=job, periodic_id=periodic_id, defer_timestamp=defer_timestamp ) + if tb_id is not None and job_id is not None: + update_task_safe(tb_id, external_id=str(job_id)) + return job_id jm.defer_periodic_job = patched diff --git a/tests/test_celery.py b/tests/test_celery.py index ac53bb2..c1b8987 100644 --- a/tests/test_celery.py +++ b/tests/test_celery.py @@ -98,7 +98,13 @@ def add_with_task_args(self, a, b): assert result.get(timeout=10, propagate=True) == 4 create.assert_called_once_with( - "new_name", value_max=10, data={"foo": "bar"}, tags={"bar": "baz"}, status=StatusEnum.PENDING, queue="celery" + "new_name", + value_max=10, + data={"foo": "bar"}, + tags={"bar": "baz"}, + status=StatusEnum.PENDING, + queue="celery", + external_id=mock.ANY, ) @@ -128,7 +134,9 @@ def add_with_task_args(self, a, b): ) assert result.get(timeout=10, propagate=True) == 4 - create.assert_called_once_with("new_name", value_max=10, actions=actions, status=StatusEnum.PENDING, queue="celery") + create.assert_called_once_with( + "new_name", value_max=10, actions=actions, status=StatusEnum.PENDING, queue="celery", external_id=mock.ANY + ) @pytest.mark.usefixtures("_bind_settings") @@ -162,6 +170,7 @@ def add_with_task_args(self, a, b): data={"foo": "bar", "celery_task_args": [2, 2], "celery_task_kwargs": {}}, status=StatusEnum.PENDING, queue="celery", + external_id=mock.ANY, ) @@ -200,6 +209,7 @@ def add_with_task_kwargs(self, a, b, c=0): actions=actions, status=StatusEnum.PENDING, queue="celery", + external_id=mock.ANY, ) @@ -237,6 +247,7 @@ def add_task_custom_serialization(self, a): data={"celery_task_args": [{"__type__": "A", "__value__": [2, 2]}], "celery_task_kwargs": {}}, status=StatusEnum.PENDING, queue="celery", + external_id=mock.ANY, ) @@ -264,7 +275,9 @@ def add_with_task_args_in_decorator(self, a, b): result = add_with_task_args_in_decorator.delay(2, 2) assert result.get(timeout=10, propagate=True) == 4 - create.assert_called_once_with(mock.ANY, status=StatusEnum.PENDING, monitor_id="123", value_max=10, queue="celery") + create.assert_called_once_with( + mock.ANY, status=StatusEnum.PENDING, monitor_id="123", value_max=10, queue="celery", external_id=mock.ANY + ) @pytest.mark.usefixtures("_bind_settings") @@ -288,6 +301,49 @@ def add_custom_queue(self, a, b): assert create.call_args.kwargs["queue"] == "high_priority" +@pytest.mark.usefixtures("_bind_settings") +def test_celery_records_task_id_as_external_id(celery_session_app, celery_session_worker): + @celery_session_app.task(bind=True, base=Task) + def add_external_id(self, a, b): + return a + b + + celery_session_worker.reload() + + with ( + mock.patch("taskbadger.celery.create_task_safe") as create, + mock.patch("taskbadger.celery.update_task_safe"), + mock.patch("taskbadger.sdk.get_task"), + ): + create.return_value = task_for_test() + result = add_external_id.apply_async((2, 2), queue="high_priority") + + assert create.call_args.kwargs["external_id"] == result.id + + +@pytest.mark.usefixtures("_bind_settings") +def test_celery_external_id_not_cached_across_invocations(celery_session_app, celery_session_worker): + # Tasks defined with taskbadger_kwargs share a class-level dict; external_id + # must not leak from one publish to the next. + @celery_session_app.task(bind=True, base=Task, taskbadger_kwargs={"value_max": 10}) + def add_repeat(self, a, b): + return a + b + + celery_session_worker.reload() + + with ( + mock.patch("taskbadger.celery.create_task_safe") as create, + mock.patch("taskbadger.celery.update_task_safe"), + mock.patch("taskbadger.sdk.get_task"), + ): + create.return_value = task_for_test() + result1 = add_repeat.apply_async((2, 2), queue="high_priority") + result2 = add_repeat.apply_async((3, 3), queue="high_priority") + + external_ids = [c.kwargs["external_id"] for c in create.call_args_list] + assert external_ids == [result1.id, result2.id] + assert result1.id != result2.id + + @pytest.mark.usefixtures("_bind_settings") def test_celery_task_retry(celery_session_app, celery_session_worker): """Note: When a task is retried, the celery task ID remains the same but a new TB task diff --git a/tests/test_celery_system_integration.py b/tests/test_celery_system_integration.py index 4944a20..b244819 100644 --- a/tests/test_celery_system_integration.py +++ b/tests/test_celery_system_integration.py @@ -123,6 +123,7 @@ def add_normal(self, a, b): status=StatusEnum.PENDING, data={"celery_task_args": [2, 2], "celery_task_kwargs": {}}, queue="celery", + external_id=mock.ANY, ) assert get_task.call_count == 1 assert update.call_count == 2 @@ -154,7 +155,10 @@ def add_normal_with_override(a, b): assert result.get(timeout=10, propagate=True) == 4 create.assert_called_once_with( - "tests.test_celery_system_integration.add_normal_with_override", status=StatusEnum.PENDING, queue="celery" + "tests.test_celery_system_integration.add_normal_with_override", + status=StatusEnum.PENDING, + queue="celery", + external_id=mock.ANY, ) @@ -189,6 +193,7 @@ def add_with_tags(a, b): status=StatusEnum.PENDING, tags=TaskRequestTags.from_dict({"tag1": "value1", "tag2": "override"}), queue="celery", + external_id=result.id, ) create.assert_called_with( client=mock.ANY, diff --git a/tests/test_procrastinate.py b/tests/test_procrastinate.py index ca9223a..7854261 100644 --- a/tests/test_procrastinate.py +++ b/tests/test_procrastinate.py @@ -115,7 +115,10 @@ def add3(a, b): _instrument_task(add3, system=None, manual=True) tb = task_for_test() - with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create: + with ( + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create, + mock.patch("taskbadger.procrastinate.update_task_safe"), + ): add3.defer(a=1, b=2) create.assert_called_once() @@ -137,7 +140,10 @@ def add_queued(a, b): _instrument_task(add_queued, system=None, manual=True) tb = task_for_test() - with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create: + with ( + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create, + mock.patch("taskbadger.procrastinate.update_task_safe"), + ): add_queued.defer(a=1, b=2) assert create.call_args.kwargs["queue"] == "high_priority" @@ -168,13 +174,69 @@ async def add5(a, b): _instrument_task(add5, system=None, manual=True) tb = task_for_test() - with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb): + with ( + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb), + mock.patch("taskbadger.procrastinate.update_task_safe"), + ): asyncio.run(add5.defer_async(a=1, b=2)) jobs = list(app.connector.jobs.values()) assert jobs[0]["args"][TB_TASK_ID_KWARG] == tb.id +@pytest.mark.usefixtures("_bind_settings") +def test_defer_records_job_id_as_external_id(app): + @app.task(name="add_ext") + def add_ext(a, b): + return a + b + + _instrument_task(add_ext, system=None, manual=True) + + tb = task_for_test() + with ( + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb), + mock.patch("taskbadger.procrastinate.update_task_safe") as update, + ): + job_id = add_ext.defer(a=1, b=2) + + update.assert_called_once_with(tb.id, external_id=str(job_id)) + + +@pytest.mark.usefixtures("_bind_settings") +def test_defer_async_records_job_id_as_external_id(app): + @app.task(name="add_ext_async") + async def add_ext_async(a, b): + return a + b + + _instrument_task(add_ext_async, system=None, manual=True) + + tb = task_for_test() + with ( + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb), + mock.patch("taskbadger.procrastinate.update_task_safe") as update, + ): + job_id = asyncio.run(add_ext_async.defer_async(a=1, b=2)) + + update.assert_called_once_with(tb.id, external_id=str(job_id)) + + +def test_defer_no_external_id_when_untracked(app): + @app.task(name="add_untracked") + def add_untracked(a, b): + return a + b + + _instrument_task(add_untracked, system=None, manual=True) + + # Badger is not configured, so no pending task is created and nothing to update. + with ( + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=None), + mock.patch("taskbadger.procrastinate.update_task_safe") as update, + ): + add_untracked.defer(a=1, b=2) + + update.assert_not_called() + + @pytest.mark.usefixtures("_bind_settings") def test_end_to_end_via_worker(app): @app.task(name="add6") @@ -194,7 +256,7 @@ def add6(a, b): app.run_worker(wait=False, install_signal_handlers=False, listen_notify=False) create.assert_called_once() - statuses = [c.kwargs["status"] for c in update.call_args_list] + statuses = [c.kwargs["status"] for c in update.call_args_list if "status" in c.kwargs] assert statuses == [StatusEnum.PROCESSING, StatusEnum.SUCCESS] @@ -206,7 +268,10 @@ def bare(a): return a tb = task_for_test() - with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb): + with ( + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb), + mock.patch("taskbadger.procrastinate.update_task_safe"), + ): bare.defer(a=1) assert getattr(bare, "_taskbadger_manual") is True @@ -223,7 +288,10 @@ def raw(a): return a tb = task_for_test() - with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create: + with ( + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create, + mock.patch("taskbadger.procrastinate.update_task_safe"), + ): raw.defer(a=1) create.assert_called_once() @@ -248,7 +316,10 @@ def dup(a): # Two @track applications must not double-wrap; defer once still creates one # PENDING task and injects one id. tb = task_for_test() - with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create: + with ( + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create, + mock.patch("taskbadger.procrastinate.update_task_safe"), + ): dup.defer(a=1) assert create.call_count == 1 jobs = list(app.connector.jobs.values()) @@ -312,7 +383,7 @@ def self_complete(): # The wrapper's post-call SUCCESS update is skipped because the cached # task is already SUCCESS. PROCESSING update is still allowed (early path). - statuses = [c.kwargs["status"] for c in update.call_args_list] + statuses = [c.kwargs["status"] for c in update.call_args_list if "status" in c.kwargs] assert StatusEnum.PROCESSING in statuses # Last attempted SUCCESS call should be suppressed assert statuses.count(StatusEnum.SUCCESS) == 0 @@ -326,7 +397,10 @@ def recorder(a, b): return a + b tb = task_for_test() - with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create: + with ( + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create, + mock.patch("taskbadger.procrastinate.update_task_safe"), + ): recorder.defer(a=5, b=6) assert create.call_args.kwargs["data"] == {