Merge pull request #57 from taskbadger/sk/ctask-none

Guard against unregistered task in before_task_publish handler

Author: snopoke

taskbadger/celery.py

@@ -170,7 +170,8 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
     if task:
         meta = {TB_TASK_ID: task.id}
         headers.update(meta)
-        ctask.update_state(task_id=headers["id"], state="PENDING", meta=meta)
+        if ctask:
+            ctask.update_state(task_id=headers["id"], state="PENDING", meta=meta)
 
 
 def _maybe_create_task(signal_sender):

tests/test_celery.py

@@ -17,7 +17,7 @@
 from kombu.utils.json import register_type
 
 from taskbadger import Action, EmailIntegration, StatusEnum
-from taskbadger.celery import Task
+from taskbadger.celery import Task, task_publish_handler
 from taskbadger.mug import Badger
 from tests.utils import task_for_test
 
@@ -280,6 +280,23 @@ def add_with_task_args_in_decorator(self, a, b):
     )
 
 
+@pytest.mark.usefixtures("_bind_settings")
+def test_celery_publish_handler_task_not_registered_locally():
+    """before_task_publish can fire in a process that never imported the task
+    being sent (e.g. a producer-only service), so `celery.current_app.tasks.get`
+    returns None. The handler should still create the TaskBadger task without
+    blowing up trying to call `.update_state()` on the missing task class."""
+
+    with mock.patch("taskbadger.celery.create_task_safe") as create:
+        create.return_value = task_for_test()
+
+        headers = {"id": "abc123", "task": "unregistered.task", "taskbadger_track": True}
+        task_publish_handler(sender="unregistered.task", headers=headers, body=[[], {}, {}])
+
+    create.assert_called_once()
+    assert headers["taskbadger_task_id"] == create.return_value.id
+
+
 @pytest.mark.usefixtures("_bind_settings")
 def test_celery_task_custom_queue(celery_session_app, celery_session_worker):
     @celery_session_app.task(bind=True, base=Task)