Add integration tests for context provider wiring

Covers the Celery, Procrastinate, @track, and Task.error() entry points with a real provider, not just the low-level helpers.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: snopoke

tests/test_celery_error.py

@@ -3,7 +3,16 @@
 import pytest
 
 from taskbadger import StatusEnum
-from taskbadger.celery import Task
+from taskbadger._error_context import capture_error_data
+from taskbadger.celery import (
+    TB_ERROR_CTX_TOKEN,
+    TB_TASK_ID,
+    Task,
+    task_failure_handler,
+    task_postrun_handler,
+    task_prerun_handler,
+)
+from taskbadger.context_providers import ContextProvider
 from taskbadger.mug import Badger
 from tests.utils import task_for_test
 
@@ -41,3 +50,103 @@ def add_error(self, a, b):
     data_kwarg = update.call_args_list[1][1]["data"]
     assert "Traceback" in data_kwarg["exception"]
     assert Badger.current.session().client is None
+
+
+class _FakeRequest(dict):
+    """Minimal stand-in for Celery's request `Context`: dict-like plus a
+    `.headers` attribute, which is all `_get_taskbadger_task_id` needs."""
+
+    headers = None
+
+
+class _FakeEinfo:
+    """Minimal stand-in for Celery's `ExceptionInfo`: wraps the exception and
+    renders a traceback-shaped string, mirroring what `task_failure`/`task_retry`
+    signals actually pass to `_update_task`."""
+
+    def __init__(self, exc):
+        self.exception = exc
+
+    def __str__(self):
+        return f"Traceback (most recent call last):\n{self.exception!r}"
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_signal_handlers_wire_context_provider_snapshot_to_failure():
+    """`task_prerun_handler` snapshots providers and stashes the token on the
+    request; `task_failure_handler` should see that exact snapshot when
+    building error data; `task_postrun_handler` then resets it."""
+
+    seen_snapshots = []
+
+    class TrackingProvider(ContextProvider):
+        identifier = "tracking"
+
+        def snapshot(self):
+            return "baseline"
+
+        def capture_error_context(self, exception, snapshot=None):
+            seen_snapshots.append(snapshot)
+            return {"snapshot": snapshot}
+
+    Badger.current.settings.context_providers = [TrackingProvider()]
+
+    sender = mock.Mock()
+    sender.request = _FakeRequest({TB_TASK_ID: "tb-1"})
+    task = task_for_test(id="tb-1", status=StatusEnum.PROCESSING)
+    sender.taskbadger_task = task
+
+    with (
+        mock.patch("taskbadger.celery.safe_get_task", return_value=task),
+        mock.patch("taskbadger.celery.update_task_safe", return_value=task) as update,
+        mock.patch("taskbadger.celery.enter_session"),
+        mock.patch("taskbadger.celery.exit_session"),
+    ):
+        task_prerun_handler(sender=sender)
+        assert TB_ERROR_CTX_TOKEN in sender.request
+
+        task_failure_handler(sender=sender, einfo=_FakeEinfo(ValueError("boom")))
+        task_postrun_handler(sender=sender)
+
+    # The snapshot seen at failure time is the one taken at prerun, not a
+    # missing/None one -- proving the token round-trips through the request.
+    assert seen_snapshots == ["baseline"]
+    data_kwarg = update.call_args.kwargs["data"]
+    assert data_kwarg["tracking"] == {"snapshot": "baseline"}
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_postrun_resets_context_after_error():
+    """Once `task_postrun_handler` runs, a later error in the same thread with
+    no provider snapshot taken shouldn't see a stale one left over from the
+    previous task."""
+
+    class TrackingProvider(ContextProvider):
+        identifier = "tracking"
+
+        def snapshot(self):
+            return "baseline"
+
+        def capture_error_context(self, exception, snapshot=None):
+            return {"snapshot": snapshot}
+
+    Badger.current.settings.context_providers = [TrackingProvider()]
+
+    task = task_for_test(id="tb-2", status=StatusEnum.PROCESSING)
+    sender = mock.Mock()
+    sender.request = _FakeRequest({TB_TASK_ID: "tb-2"})
+    sender.taskbadger_task = task
+
+    with (
+        mock.patch("taskbadger.celery.safe_get_task", return_value=task),
+        mock.patch("taskbadger.celery.update_task_safe", return_value=task),
+        mock.patch("taskbadger.celery.enter_session"),
+        mock.patch("taskbadger.celery.exit_session"),
+    ):
+        task_prerun_handler(sender=sender)
+        task_postrun_handler(sender=sender)
+
+    # A failure reported outside of any tracked task's prerun/postrun window
+    # (e.g. directly via Task.error) has no snapshot to compare against.
+    data = capture_error_data(ValueError("boom"))
+    assert data["tracking"] == {"snapshot": None}

tests/test_decorators.py

@@ -1,6 +1,7 @@
 from unittest import mock
 
 from taskbadger import track
+from taskbadger.context_providers import ContextProvider
 from taskbadger.mug import Badger, Settings
 
 
@@ -52,6 +53,32 @@ def test(arg):
     assert update.call_args.kwargs["data"]["exception"] == "test"
 
 
+@mock.patch("taskbadger.decorators.create_task_safe")
+@mock.patch("taskbadger.decorators._update_safe")
+def test_track_decorator_error_consults_context_provider(update, create):
+    class FakeProvider(ContextProvider):
+        identifier = "fake"
+
+        def capture_error_context(self, exception, snapshot=None):
+            return {"detail": str(exception)}
+
+    Badger.current.bind(Settings("https://taskbadger.net", "token", "org", "proj", context_providers=[FakeProvider()]))
+    try:
+
+        @track
+        def test(arg):
+            raise Exception("test")
+
+        try:
+            test("test")
+        except Exception:
+            pass
+    finally:
+        Badger.current.bind(None)
+
+    assert update.call_args.kwargs["data"] == {"exception": "test", "fake": {"detail": "test"}}
+
+
 @mock.patch("taskbadger.decorators._update_safe")
 def test_track_decorator_badger_not_configured(update):
     @track

tests/test_procrastinate.py

@@ -7,6 +7,8 @@
 from procrastinate import testing
 
 from taskbadger import StatusEnum
+from taskbadger.context_providers import ContextProvider
+from taskbadger.mug import Badger
 from taskbadger.procrastinate import TB_TASK_ID_KWARG, _instrument_task, current_task, track
 from tests.utils import task_for_test
 
@@ -91,6 +93,35 @@ def boom():
     assert err_call.kwargs["data"] == {"x": 1, "exception": "nope"}
 
 
+@pytest.mark.usefixtures("_bind_settings")
+def test_worker_marks_error_consults_context_provider(app):
+    class FakeProvider(ContextProvider):
+        identifier = "fake"
+
+        def capture_error_context(self, exception, snapshot=None):
+            return {"detail": str(exception)}
+
+    Badger.current.settings.context_providers = [FakeProvider()]
+
+    @app.task(name="boom_with_provider")
+    def boom():
+        raise ValueError("nope")
+
+    _instrument_task(boom, system=None, manual=True)
+
+    with (
+        mock.patch("taskbadger.procrastinate.update_task_safe") as update,
+        mock.patch("taskbadger.sdk.get_task") as get,
+    ):
+        get.return_value = task_for_test(status=StatusEnum.PROCESSING)
+        update.return_value = task_for_test(status=StatusEnum.PROCESSING)
+        with pytest.raises(ValueError, match="nope"):
+            boom.func(**{TB_TASK_ID_KWARG: "tb-provider"})
+
+    err_call = update.call_args_list[-1]
+    assert err_call.kwargs["data"] == {"exception": "nope", "fake": {"detail": "nope"}}
+
+
 @pytest.mark.usefixtures("_bind_settings")
 def test_worker_no_id_runs_clean(app):
     @app.task(name="add2")

tests/test_sdk.py

@@ -6,6 +6,7 @@
 import pytest
 
 from taskbadger import Action, EmailIntegration, StatusEnum, WebhookIntegration, create_task
+from taskbadger.context_providers import ContextProvider
 from taskbadger.exceptions import TaskbadgerException
 from taskbadger.internal.models import (
     PatchedTaskRequest,
@@ -207,6 +208,63 @@ def test_update_data(settings, patched_update):
     _verify_update(settings, patched_update, data={"a": 1})
 
 
+def test_error_with_exception(settings, patched_update):
+    api_task = task_for_test()
+    task = Task(api_task)
+
+    patched_update.return_value = Response(HTTPStatus.OK, b"", {}, api_task)
+    task.error(exception=ValueError("boom"))
+
+    _verify_update(settings, patched_update, status=StatusEnum.ERROR, data={"exception": "boom"})
+
+
+def test_error_with_exception_consults_context_provider(settings, patched_update):
+    class FakeProvider(ContextProvider):
+        identifier = "fake"
+
+        def capture_error_context(self, exception, snapshot=None):
+            return {"detail": str(exception)}
+
+    settings.context_providers = [FakeProvider()]
+
+    api_task = task_for_test()
+    task = Task(api_task)
+    patched_update.return_value = Response(HTTPStatus.OK, b"", {}, api_task)
+    task.error(exception=ValueError("boom"))
+
+    _verify_update(
+        settings,
+        patched_update,
+        status=StatusEnum.ERROR,
+        data={"exception": "boom", "fake": {"detail": "boom"}},
+    )
+
+
+def test_error_explicit_data_overrides_provider_data(settings, patched_update):
+    """Explicit `data` passed to `error()` wins over provider-derived data for
+    overlapping keys."""
+
+    class FakeProvider(ContextProvider):
+        identifier = "fake"
+
+        def capture_error_context(self, exception, snapshot=None):
+            return {"detail": "from provider"}
+
+    settings.context_providers = [FakeProvider()]
+
+    api_task = task_for_test()
+    task = Task(api_task)
+    patched_update.return_value = Response(HTTPStatus.OK, b"", {}, api_task)
+    task.error(exception=ValueError("boom"), data={"exception": "overridden", "extra": 1})
+
+    _verify_update(
+        settings,
+        patched_update,
+        status=StatusEnum.ERROR,
+        data={"exception": "overridden", "extra": 1, "fake": {"detail": "from provider"}},
+    )
+
+
 def test_increment_value(settings, patched_update):
     api_task = task_for_test()
     task = Task(api_task)