feat(procrastinate): @track decorator with bare and parameterized forms

Author: snopoke

taskbadger/procrastinate.py

@@ -222,3 +222,41 @@ def _maybe_create_pending(task, kwargs):
     new_kwargs = dict(kwargs)
     new_kwargs[TB_TASK_ID_KWARG] = tb_task.id
     return new_kwargs
+
+
+_TRACK_OPT_KEYS = ("name", "value_max", "tags", "data", "record_task_args")
+
+
+def track(original_task=None, **opts):
+    """Opt a Procrastinate task into TaskBadger tracking.
+
+    Usage:
+
+        @track
+        @app.task(...)
+        def my_task(...): ...
+
+        @track(name="custom", value_max=100, tags={"env": "prod"})
+        @app.task(...)
+        async def big_job(...): ...
+
+    Accepted keyword options (all optional):
+        name: TaskBadger task name (defaults to the Procrastinate task's name).
+        value_max: Maximum value for the TaskBadger task.
+        tags: Dict of tags applied to the TaskBadger task.
+        data: Dict of initial data merged into the TaskBadger task.
+        record_task_args: If True, serialize the Procrastinate job kwargs and
+            store them under ``data["procrastinate_task_kwargs"]``. Defaults to
+            ``None`` meaning "inherit from system integration if any, else False".
+    """
+    unknown = set(opts) - set(_TRACK_OPT_KEYS)
+    if unknown:
+        raise TypeError(f"track() got unexpected keyword arguments: {sorted(unknown)}")
+
+    def wrap(task):
+        _instrument_task(task, system=None, manual=True, opts=opts)
+        return task
+
+    if original_task is None:
+        return wrap
+    return wrap(original_task)

tests/test_procrastinate.py

@@ -7,7 +7,7 @@
 from procrastinate import testing
 
 from taskbadger import StatusEnum
-from taskbadger.procrastinate import TB_TASK_ID_KWARG, _instrument_task
+from taskbadger.procrastinate import TB_TASK_ID_KWARG, _instrument_task, track
 from tests.utils import task_for_test
 
 
@@ -181,3 +181,70 @@ def add6(a, b):
     create.assert_called_once()
     statuses = [c.kwargs["status"] for c in update.call_args_list]
     assert statuses == [StatusEnum.PROCESSING, StatusEnum.SUCCESS]
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_track_bare_form(app):
+    @track
+    @app.task(name="bare")
+    def bare(a):
+        return a
+
+    tb = task_for_test()
+    with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb):
+        bare.defer(a=1)
+
+    assert getattr(bare, "_taskbadger_manual") is True
+    # Inspect the actual Procrastinate job - jobs is a dict keyed by int, kwargs under "args"
+    jobs = list(app.connector.jobs.values())
+    assert jobs[0]["args"][TB_TASK_ID_KWARG] == tb.id
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_track_parameterized(app):
+    @track(name="custom", value_max=10, tags={"env": "test"}, data={"k": "v"})
+    @app.task(name="raw_name")
+    def raw(a):
+        return a
+
+    tb = task_for_test()
+    with mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create:
+        raw.defer(a=1)
+
+    create.assert_called_once()
+    assert create.call_args.args == ("custom",)
+    assert create.call_args.kwargs == {
+        "status": StatusEnum.PENDING,
+        "value_max": 10,
+        "tags": {"env": "test"},
+        "data": {"k": "v"},
+    }
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_track_idempotent(app):
+    @track
+    @track
+    @app.task(name="dup")
+    def dup(a):
+        return 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:
+        dup.defer(a=1)
+    assert create.call_count == 1
+    jobs = list(app.connector.jobs.values())
+    args = jobs[0]["args"]
+    # Reserved key appears exactly once
+    assert list(args).count(TB_TASK_ID_KWARG) == 1
+
+
+def test_track_unknown_opt_raises(app):
+    @app.task(name="bad")
+    def bad():
+        pass
+
+    with pytest.raises(TypeError, match="unexpected keyword"):
+        track(name="x", does_not_exist=True)(bad)