run celery with redis backend and fix serialization issues

Author: snopoke

.github/workflows/tests.yml

@@ -8,12 +8,27 @@ on:
 jobs:
   python-tests:
     runs-on: ubuntu-latest
+    services:
+      redis:
+        image: redis
+        options: >-
+          --health-cmd "redis-cli ping"
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+        ports:
+          - 6379:6379
     strategy:
       max-parallel: 4
       matrix:
         python-version:
           - "3.9"
           - "3.10"
+          - "3.11"
+          - "3.12"
+        celery-version:
+          - ">=5.3,<5.4"
+          - ">=5.4"
     steps:
       - uses: actions/checkout@v3
       - name: Set up Python ${{ matrix.python-version }}
@@ -31,6 +46,8 @@ jobs:
           poetry --version
           poetry check --no-interaction
       - name: Install project
-        run: poetry install --no-interaction
+        run: |
+          poetry install --no-interaction
+          pip install celery"${{ matrix.celery-version }}"
       - name: Run tests
         run: poetry run pytest -v

taskbadger/celery.py

@@ -156,6 +156,7 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
     celery_system = Badger.current.settings.get_system_by_id("celery")
     auto_track = celery_system and celery_system.track_task(sender)
     manual_track = headers.get("taskbadger_track")
+    header_kwargs = headers.pop(TB_KWARGS_ARG, {})
     if not manual_track and not auto_track:
         return
 
@@ -168,7 +169,7 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
             kwargs[attr.removeprefix(KWARG_PREFIX)] = getattr(ctask, attr)
 
     # get kwargs from the task headers (set via apply_async)
-    kwargs.update(headers.get(TB_KWARGS_ARG, {}))
+    kwargs.update(header_kwargs)
     kwargs["status"] = StatusEnum.PENDING
     name = kwargs.pop("name", headers["task"])
 

tests/conftest.py

@@ -8,3 +8,12 @@ def bind_settings():
     Badger.current.bind(Settings("https://taskbadger.net", "token", "org", "proj"))
     yield
     Badger.current.bind(None)
+
+
+@pytest.fixture(scope="session", autouse=True)
+def celery_config():
+    """Test against Redis to ensure serialization works"""
+    return {
+        "broker_url": "redis://localhost:6379",
+        "result_backend": "redis://localhost:6379",
+    }

tests/test_celery_error.py

@@ -21,6 +21,7 @@ def add_error(self, a, b):
         "taskbadger.celery.update_task_safe"
     ) as update, mock.patch("taskbadger.celery.get_task") as get_task:
         task = task_for_test()
+        create.return_value = task
         get_task.return_value = task
         update.return_value = task
         result = add_error.delay(2, 2)