refactor: drop get_task_fn parameter from shared safe_get_task

Author: snopoke

taskbadger/_integrations.py

@@ -2,11 +2,9 @@
 (Celery, Procrastinate). Not part of the public API.
 
 Each integration creates its own module-level ``TaskCache`` instance and
-defines a thin ``safe_get_task`` wrapper that reads ``get_task`` from the
-integration module's own globals (so existing test mocks on
-``taskbadger.celery.get_task`` / ``taskbadger.procrastinate.get_task`` keep
-working). ``BaseSystemIntegration`` provides the common ctor/include-exclude
-shape; subclasses override ``track_task`` if they need to filter additional
+defines a thin ``safe_get_task`` wrapper around the shared one defined here.
+``BaseSystemIntegration`` provides the common ctor/include-exclude shape;
+subclasses override ``track_task`` if they need to filter additional
 task names (e.g. Procrastinate built-ins).
 """
 
@@ -15,8 +13,8 @@
 import collections
 import logging
 import re
-from collections.abc import Callable
 
+from . import sdk
 from .internal.models import StatusEnum
 from .systems import System
 
@@ -53,21 +51,16 @@ def unset(self, key) -> None:
         self.cache.pop(key, None)
 
 
-def safe_get_task(cache: TaskCache, task_id: str, get_task_fn: Callable):
+def safe_get_task(cache: TaskCache, task_id: str):
     """Cache-aware ``get_task``: returns the cached entry if present, otherwise
-    fetches via ``get_task_fn`` and caches the result. Errors are logged and
+    fetches via the SDK ``get_task`` and caches the result. Errors are logged and
     swallowed (returns ``None``). ``None`` results are not cached.
-
-    ``get_task_fn`` is passed in (rather than imported here) so callers can
-    use their own module-level ``get_task`` reference — this keeps existing
-    test patches on ``taskbadger.celery.get_task`` / ``taskbadger.procrastinate.get_task``
-    intercepting the fetch.
     """
     cached = cache.get(task_id)
     if cached is not None:
         return cached
     try:
-        task = get_task_fn(task_id)
+        task = sdk.get_task(task_id)
     except Exception as e:
         log.warning("Error fetching task '%s': %s", task_id, e)
         return None

taskbadger/celery.py

@@ -12,12 +12,13 @@
 )
 from kombu import serialization
 
+from . import sdk
 from ._integrations import TERMINAL_STATES, TaskCache
 from ._integrations import safe_get_task as _shared_safe_get_task
 from .internal.models import StatusEnum
 from .mug import Badger
 from .safe_sdk import create_task_safe, update_task_safe
-from .sdk import DefaultMergeStrategy, get_task
+from .sdk import DefaultMergeStrategy
 
 KWARG_PREFIX = "taskbadger_"
 TB_KWARGS_ARG = f"{KWARG_PREFIX}kwargs"
@@ -91,7 +92,7 @@ def apply_async(self, *args, **kwargs):
         tb_task_id = info.get(TB_TASK_ID) if isinstance(info, dict) else None
         setattr(result, TB_TASK_ID, tb_task_id)
 
-        _get_task = functools.partial(get_task, tb_task_id) if tb_task_id else lambda: None
+        _get_task = functools.partial(sdk.get_task, tb_task_id) if tb_task_id else lambda: None
         setattr(result, "get_taskbadger_task", _get_task)
 
         return result
@@ -328,7 +329,7 @@ def exit_session(signal_sender):
 
 
 def safe_get_task(task_id: str):
-    return _shared_safe_get_task(_task_cache, task_id, get_task)
+    return _shared_safe_get_task(_task_cache, task_id)
 
 
 def _get_taskbadger_task_id(request):

taskbadger/procrastinate.py

@@ -22,7 +22,7 @@
 from .internal.models import StatusEnum
 from .mug import Badger
 from .safe_sdk import create_task_safe, update_task_safe
-from .sdk import DefaultMergeStrategy, get_task
+from .sdk import DefaultMergeStrategy
 
 log = logging.getLogger("taskbadger")
 
@@ -141,7 +141,7 @@ def _update_status(tb_id, status, exception=None):
 
 
 def _safe_get_task(task_id):
-    return _shared_safe_get_task(_task_cache, task_id, get_task)
+    return _shared_safe_get_task(_task_cache, task_id)
 
 
 def _wrap_defer(task):

tests/test_celery.py

@@ -43,7 +43,7 @@ def add_normal(self, a, b):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe") as update,
-        mock.patch("taskbadger.celery.get_task") as get_task,
+        mock.patch("taskbadger.sdk.get_task") as get_task,
     ):
         tb_task = task_for_test()
         create.return_value = tb_task
@@ -71,7 +71,7 @@ def add_with_task_args(self, a, b):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe"),
-        mock.patch("taskbadger.celery.get_task"),
+        mock.patch("taskbadger.sdk.get_task"),
     ):
         create.return_value = task_for_test()
 
@@ -100,7 +100,7 @@ def add_with_task_args(self, a, b):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe"),
-        mock.patch("taskbadger.celery.get_task"),
+        mock.patch("taskbadger.sdk.get_task"),
     ):
         create.return_value = task_for_test()
 
@@ -129,7 +129,7 @@ def add_with_task_args(self, a, b):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe"),
-        mock.patch("taskbadger.celery.get_task"),
+        mock.patch("taskbadger.sdk.get_task"),
     ):
         create.return_value = task_for_test()
 
@@ -162,7 +162,7 @@ def add_with_task_kwargs(self, a, b, c=0):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe"),
-        mock.patch("taskbadger.celery.get_task"),
+        mock.patch("taskbadger.sdk.get_task"),
     ):
         create.return_value = task_for_test()
 
@@ -206,7 +206,7 @@ def add_task_custom_serialization(self, a):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe"),
-        mock.patch("taskbadger.celery.get_task"),
+        mock.patch("taskbadger.sdk.get_task"),
     ):
         create.return_value = task_for_test()
 
@@ -240,7 +240,7 @@ def add_with_task_args_in_decorator(self, a, b):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe"),
-        mock.patch("taskbadger.celery.get_task"),
+        mock.patch("taskbadger.sdk.get_task"),
     ):
         create.return_value = task_for_test()
 
@@ -272,7 +272,7 @@ def add_retry(self, a, b, is_retry=False):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe") as update,
-        mock.patch("taskbadger.celery.get_task") as get_task,
+        mock.patch("taskbadger.sdk.get_task") as get_task,
     ):
         create.return_value = task_for_test()
         get_task.return_value = task_for_test()
@@ -292,7 +292,7 @@ def add_retry(self, a, b, is_retry=False):
 def test_celery_task_badger_not_configured(celery_session_app, celery_session_worker):
     @celery_session_app.task(bind=True, base=Task)
     def add_no_tb(self, a, b):
-        with mock.patch("taskbadger.celery.get_task") as get_task:
+        with mock.patch("taskbadger.sdk.get_task") as get_task:
             assert self.taskbadger_task_id is None
             assert Badger.current.session().client is None
             get_task.assert_not_called()
@@ -323,7 +323,7 @@ def add_no_tb(self, a, b):
 def test_task_direct_call(celery_session_app, celery_session_worker):
     @celery_session_app.task(bind=True, base=Task)
     def add_direct(self, a, b):
-        with mock.patch("taskbadger.celery.get_task") as get_task:
+        with mock.patch("taskbadger.sdk.get_task") as get_task:
             assert self.taskbadger_task_id is None
             assert Badger.current.session().client is None
             get_task.assert_not_called()
@@ -351,7 +351,7 @@ def add_shared_task(self, a, b):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe") as update,
-        mock.patch("taskbadger.celery.get_task"),
+        mock.patch("taskbadger.sdk.get_task"),
     ):
         create.return_value = task_for_test()
 
@@ -378,7 +378,7 @@ def task_signature(self, a):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe") as update,
-        mock.patch("taskbadger.celery.get_task") as get_task,
+        mock.patch("taskbadger.sdk.get_task") as get_task,
     ):
         create.return_value = task_for_test()
 
@@ -410,7 +410,7 @@ def task_map_fn(self, a):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe") as update,
-        mock.patch("taskbadger.celery.get_task"),
+        mock.patch("taskbadger.sdk.get_task"),
     ):
         tb_task = task_for_test()
         create.return_value = tb_task
@@ -443,7 +443,7 @@ def task_starmap_fn(self, a, b):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe") as update,
-        mock.patch("taskbadger.celery.get_task"),
+        mock.patch("taskbadger.sdk.get_task"),
     ):
         tb_task = task_for_test()
         create.return_value = tb_task
@@ -476,7 +476,7 @@ def task_chunks_fn(self, a):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe") as update,
-        mock.patch("taskbadger.celery.get_task"),
+        mock.patch("taskbadger.sdk.get_task"),
     ):
         tb_task = task_for_test()
         create.return_value = tb_task
@@ -506,7 +506,7 @@ def add_manual_update(self, a, b, is_retry=False):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe") as update,
-        mock.patch("taskbadger.celery.get_task") as get_task,
+        mock.patch("taskbadger.sdk.get_task") as get_task,
     ):
         create.return_value = task_for_test()
 

tests/test_celery_error.py

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

tests/test_celery_system_integration.py

@@ -77,7 +77,7 @@ def add_normal(self, a, b):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe") as update,
-        mock.patch("taskbadger.celery.get_task") as get_task,
+        mock.patch("taskbadger.sdk.get_task") as get_task,
     ):
         tb_task = task_for_test()
         create.return_value = tb_task
@@ -109,7 +109,7 @@ def add_normal(self, a, b):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe") as update,
-        mock.patch("taskbadger.celery.get_task") as get_task,
+        mock.patch("taskbadger.sdk.get_task") as get_task,
     ):
         tb_task = task_for_test()
         create.return_value = tb_task
@@ -144,7 +144,7 @@ def add_normal_with_override(a, b):
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe"),
-        mock.patch("taskbadger.celery.get_task"),
+        mock.patch("taskbadger.sdk.get_task"),
     ):
         tb_task = task_for_test()
         create.return_value = tb_task
@@ -169,7 +169,7 @@ def add_with_tags(a, b):
     with (
         mock.patch("taskbadger.sdk.task_create.sync_detailed") as create,
         mock.patch("taskbadger.celery.update_task_safe"),
-        mock.patch("taskbadger.celery.get_task"),
+        mock.patch("taskbadger.sdk.get_task"),
     ):
         tb_task = task_for_test()
         create.return_value = Response(

tests/test_procrastinate.py

@@ -44,7 +44,7 @@ def add(a, b):
 
     with (
         mock.patch("taskbadger.procrastinate.update_task_safe") as update,
-        mock.patch("taskbadger.procrastinate.get_task") as get,
+        mock.patch("taskbadger.sdk.get_task") as get,
     ):
         get.return_value = task_for_test(status=StatusEnum.PROCESSING)
         add.func(a=2, b=3, **{TB_TASK_ID_KWARG: "tb-123"})
@@ -65,7 +65,7 @@ async def add_async(a, b):
 
     with (
         mock.patch("taskbadger.procrastinate.update_task_safe") as update,
-        mock.patch("taskbadger.procrastinate.get_task") as get,
+        mock.patch("taskbadger.sdk.get_task") as get,
     ):
         get.return_value = task_for_test(status=StatusEnum.PROCESSING)
         result = asyncio.run(add_async.func(a=2, b=3, **{TB_TASK_ID_KWARG: "tb-456"}))
@@ -85,7 +85,7 @@ def boom():
 
     with (
         mock.patch("taskbadger.procrastinate.update_task_safe") as update,
-        mock.patch("taskbadger.procrastinate.get_task") as get,
+        mock.patch("taskbadger.sdk.get_task") as get,
     ):
         get.return_value = task_for_test(status=StatusEnum.PROCESSING, data={"x": 1})
         update.return_value = task_for_test(status=StatusEnum.PROCESSING, data={"x": 1})
@@ -179,7 +179,7 @@ def add6(a, b):
     with (
         mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create,
         mock.patch("taskbadger.procrastinate.update_task_safe") as update,
-        mock.patch("taskbadger.procrastinate.get_task") as get,
+        mock.patch("taskbadger.sdk.get_task") as get,
     ):
         get.return_value = task_for_test(id=tb.id, status=StatusEnum.PROCESSING)
         add6.defer(a=2, b=3)
@@ -270,7 +270,7 @@ def capture():
     with (
         mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb),
         mock.patch("taskbadger.procrastinate.update_task_safe", return_value=tb),
-        mock.patch("taskbadger.procrastinate.get_task", return_value=tb),
+        mock.patch("taskbadger.sdk.get_task", return_value=tb),
     ):
         capture.defer()
         app.run_worker(wait=False, install_signal_handlers=False, listen_notify=False)
@@ -296,7 +296,7 @@ def self_complete():
     with (
         mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb_pending),
         mock.patch("taskbadger.procrastinate.update_task_safe") as update,
-        mock.patch("taskbadger.procrastinate.get_task", return_value=tb_done),
+        mock.patch("taskbadger.sdk.get_task", return_value=tb_done),
     ):
         self_complete.defer()
         app.run_worker(wait=False, install_signal_handlers=False, listen_notify=False)
@@ -340,7 +340,7 @@ def ctx_task(context, a):
     with (
         mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb),
         mock.patch("taskbadger.procrastinate.update_task_safe"),
-        mock.patch("taskbadger.procrastinate.get_task", return_value=tb),
+        mock.patch("taskbadger.sdk.get_task", return_value=tb),
     ):
         ctx_task.defer(a=42)
         app.run_worker(wait=False, install_signal_handlers=False, listen_notify=False)