diff --git a/taskbadger/celery.py b/taskbadger/celery.py index e61964b..e30f65d 100644 --- a/taskbadger/celery.py +++ b/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): diff --git a/tests/test_celery.py b/tests/test_celery.py index c1b8987..e86d565 100644 --- a/tests/test_celery.py +++ b/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)