Let tracked Procrastinate tasks inherit system options

A `@track`ed task kept `_taskbadger_system = None` when the integration was
constructed after it was registered, so it never picked up the system's
options.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: snopoke

taskbadger/procrastinate.py

@@ -42,6 +42,7 @@
 _INSTRUMENTED_ATTR = "_taskbadger_instrumented"
 _MANUAL_ATTR = "_taskbadger_manual"
 _OPTS_ATTR = "_taskbadger_opts"
+_SYSTEM_ATTR = "_taskbadger_system"
 
 _current_tb_task_id: ContextVar[str | None] = ContextVar("_current_tb_task_id", default=None)
 
@@ -62,6 +63,12 @@ def _instrument_task(task, system=None, manual=False, opts=None):
     if manual:
         setattr(task, _MANUAL_ATTR, True)
 
+    # Set before the idempotency check: a task already instrumented by ``@track``
+    # when the system integration is constructed still has to inherit the
+    # system's options (heartbeat_interval, stale_timeout, record_task_args).
+    if system is not None or not hasattr(task, _SYSTEM_ATTR):
+        setattr(task, _SYSTEM_ATTR, system)
+
     if getattr(task, _INSTRUMENTED_ATTR, False):
         return
 
@@ -116,7 +123,6 @@ def wrapped(*args, **kwargs):
     _wrap_defer(task)
     task.func = wrapped
     setattr(task, _INSTRUMENTED_ATTR, True)
-    setattr(task, "_taskbadger_system", system)
 
 
 def _update_status(tb_id, status, exception=None):
@@ -150,7 +156,7 @@ def _update_status(tb_id, status, exception=None):
 def _heartbeat_interval(task):
     """Seconds between pings while ``task`` runs, or ``None`` for no heartbeat."""
     opts = getattr(task, _OPTS_ATTR, {}) or {}
-    system = getattr(task, "_taskbadger_system", None)
+    system = getattr(task, _SYSTEM_ATTR, None)
     interval, _ = resolve_heartbeat_options(opts.get("heartbeat_interval"), None, system)
     return interval
 
@@ -195,7 +201,7 @@ def _create_pending_task(task, task_kwargs, queue=None):
     if not Badger.is_configured():
         return None
 
-    system = getattr(task, "_taskbadger_system", None)
+    system = getattr(task, _SYSTEM_ATTR, None)
     manual = getattr(task, _MANUAL_ATTR, False)
     auto = bool(system) and system.track_task(task.name)
     if not manual and not auto:

tests/test_procrastinate_heartbeat.py

@@ -150,6 +150,30 @@ def auto_heartbeat(a):
     heartbeat.stop.assert_called_once_with("tb-5")
 
 
+@pytest.mark.usefixtures("_bind_settings")
+def test_system_integration_heartbeat_for_tracked_task(app, heartbeat):
+    """A ``@track``ed task still inherits the system's options when the
+    integration is constructed after the task was registered."""
+
+    @track()
+    @app.task(name="tracked_auto_heartbeat")
+    def tracked_auto_heartbeat(a):
+        return a
+
+    ProcrastinateSystemIntegration(app=app, heartbeat_interval=45)
+
+    with (
+        mock.patch("taskbadger.procrastinate.create_task_safe", return_value=task_for_test(id="tb-6")) as create,
+        mock.patch("taskbadger.procrastinate.update_task_safe"),
+        mock.patch("taskbadger.sdk.get_task", return_value=task_for_test(status=StatusEnum.PROCESSING)),
+    ):
+        tracked_auto_heartbeat.defer(a=1)
+        tracked_auto_heartbeat.func(a=1, **{TB_TASK_ID_KWARG: "tb-6"})
+
+    assert create.call_args.kwargs["stale_timeout"] == 90
+    heartbeat.start.assert_called_once_with("tb-6", 45)
+
+
 def test_track_rejects_bad_interval(app):
     with pytest.raises(ConfigurationError):
         track(heartbeat_interval=0)