feat(procrastinate): verify pass_context=True works with task wrapper

Author: snopoke

taskbadger/procrastinate.py

@@ -68,6 +68,8 @@ def _instrument_task(task, system=None, manual=False, opts=None):
     original_func = task.func
     is_async = inspect.iscoroutinefunction(original_func)
 
+    # pass_context=True works transparently: Procrastinate passes the context
+    # object as the first positional arg; our *args/**kwargs wrapper forwards it.
     if is_async:
 
         @functools.wraps(original_func)

tests/test_procrastinate.py

@@ -324,3 +324,28 @@ def recorder(a, b):
         "existing": 1,
         "procrastinate_task_kwargs": {"a": 5, "b": 6},
     }
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_pass_context_forwards_context(app):
+    seen = {}
+
+    @track
+    @app.task(name="ctx_task", pass_context=True)
+    def ctx_task(context, a):
+        seen["context"] = context
+        seen["a"] = a
+
+    tb = task_for_test()
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb),
+        mock.patch("taskbadger.procrastinate.update_task_safe"),
+        mock.patch("taskbadger.procrastinate.get_task", return_value=tb),
+    ):
+        ctx_task.defer(a=42)
+        app.run_worker(wait=False, install_signal_handlers=False, listen_notify=False)
+
+    assert seen["a"] == 42
+    # Context object should be passed through unchanged
+    assert seen["context"] is not None
+    assert seen["context"].task is ctx_task