Merge pull request #61 from taskbadger/sk/list-tasks-wrapper

Return SDK Task objects from list_tasks

Author: snopoke

taskbadger/__init__.py

@@ -3,7 +3,7 @@
 from .internal.models import StatusEnum
 from .mug import Badger, Session
 from .safe_sdk import create_task_safe, update_task_safe
-from .sdk import DefaultMergeStrategy, Task, create_task, get_task, init, list_tasks, update_task
+from .sdk import DefaultMergeStrategy, Task, TaskList, create_task, get_task, init, list_tasks, update_task
 
 __all__ = [
     "track",
@@ -17,6 +17,7 @@
     "update_task_safe",
     "DefaultMergeStrategy",
     "Task",
+    "TaskList",
     "create_task",
     "get_task",
     "init",

taskbadger/celery.py

@@ -324,18 +324,23 @@ def _maybe_create_task(signal_sender):
     delivery_info = getattr(signal_sender.request, "delivery_info", None) or {}
     queue = delivery_info.get("routing_key")
     external_id = signal_sender.request.id
+    # `before_task_publish` never ran for eager tasks, so their per-call options
+    # are still sitting in the headers rather than resolved into the message.
+    # Canvas tasks never have any here: `task_publish_handler` strips TB headers
+    # off `celery.*` messages before its early return, so nothing reaches the
+    # worker.
+    header_kwargs = headers.get(TB_KWARGS_ARG) or {}
     create_kwargs = {
         "status": StatusEnum.PENDING,
         "data": data,
         "queue": queue,
         "external_id": external_id,
         # eager and canvas tasks are created here rather than at publish time, but
-        # still run inside whatever task invoked them
-        "parent": parent_id(),
+        # still run inside whatever task invoked them. For eager tasks an explicit
+        # `taskbadger_parent` wins, as it does at publish time — including an
+        # explicit `None`, which asks for a root task.
+        "parent": header_kwargs["parent"] if "parent" in header_kwargs else parent_id(),
     }
-    # `before_task_publish` never ran for these, so per-call options are still
-    # sitting in the headers rather than resolved into the message.
-    header_kwargs = headers.get(TB_KWARGS_ARG) or {}
     heartbeat_interval, stale_timeout = resolve_heartbeat_options(
         header_kwargs.get("heartbeat_interval", getattr(signal_sender, TB_HEARTBEAT_INTERVAL, None)),
         header_kwargs.get("stale_timeout", getattr(signal_sender, TB_STALE_TIMEOUT, None)),

taskbadger/sdk.py

@@ -302,7 +302,7 @@ def update_task(
     return Task(response.parsed)
 
 
-def list_tasks(page_size: int = None, cursor: str = None, parent: str = None):
+def list_tasks(page_size: int = None, cursor: str = None, parent: str = None) -> "TaskList":
     """List tasks.
 
     Arguments:
@@ -314,7 +314,7 @@ def list_tasks(page_size: int = None, cursor: str = None, parent: str = None):
     with Session() as client:
         response = task_list.sync_detailed(client=client, **kwargs)
     _check_response(response)
-    return response.parsed
+    return TaskList(response.parsed)
 
 
 _ACTIONS_DEPRECATED_MESSAGE = (
@@ -554,6 +554,11 @@ def tags(self):
         return self._task.tags.to_dict()
 
     def __getattr__(self, item):
+        if item.startswith("_"):
+            # don't delegate private / dunder lookups: `copy` and `pickle` probe
+            # for e.g. `__setstate__` on an instance that has no `_task` yet,
+            # which would recurse until the stack blows up.
+            raise AttributeError(item)
         return getattr(self._task, item)
 
     def safe_update(self, **kwargs):
@@ -581,6 +586,39 @@ def _check_update_value_interval(self, new_value, value_step: int = None):
         return True
 
 
+class TaskList:
+    """A page of tasks as returned by [taskbadger.list_tasks][].
+
+    Iterating over a `TaskList` yields [taskbadger.Task][] objects:
+
+        for task in taskbadger.list_tasks():
+            print(task.name)
+    """
+
+    def __init__(self, task_list):
+        self._task_list = task_list
+        self._results = [Task(task) for task in task_list.results]
+
+    @property
+    def results(self) -> list[Task]:
+        """The tasks in this page."""
+        return self._results
+
+    def __iter__(self):
+        return iter(self._results)
+
+    def __len__(self):
+        return len(self._results)
+
+    def __getattr__(self, item):
+        if item.startswith("_"):
+            # don't delegate private / dunder lookups: `copy` and `pickle` probe
+            # for e.g. `__setstate__` on an instance that has no `_task_list` yet,
+            # which would recurse until the stack blows up.
+            raise AttributeError(item)
+        return getattr(self._task_list, item)
+
+
 def _none_to_unset(value):
     return UNSET if value is None else value
 

tasks.py

@@ -13,7 +13,9 @@ def tag_release(c: Context):
     if bump_key in ("1", "2", "3"):
         bump = {"1": "major", "2": "minor", "3": "patch"}[bump_key]
         version = _bump_version(bump)
-        c.run("git add pyproject.toml")
+        # the lockfile pins the project's own version, so it moves with pyproject
+        c.run("uv lock")
+        c.run("git add pyproject.toml uv.lock")
         c.run(f"git commit -m 'Bump version to {version}'")
 
     if input(f"\nReady to release version {version}? [y/n]") == "y":

tests/test_celery_system_integration.py

@@ -16,7 +16,14 @@
 from unittest import mock
 
 import pytest
-from celery.signals import task_prerun
+from celery.signals import (
+    before_task_publish,
+    task_failure,
+    task_postrun,
+    task_prerun,
+    task_retry,
+    task_success,
+)
 
 from taskbadger import StatusEnum
 from taskbadger.celery import Task
@@ -247,7 +254,22 @@ def _assert_signals(check_is_connected=True):
 
 
 def _disconnect_signals():
-    from taskbadger.celery import task_prerun_handler
+    """Disconnect every handler the module connected on import.
 
-    task_prerun.disconnect(task_prerun_handler)
+    All of them, not just the one asserted on: re-importing the module connects
+    a second copy of each handler, and anything left behind keeps firing for the
+    rest of the session — from a stale module object that `mock.patch` no longer
+    reaches.
+    """
+    import taskbadger.celery as tb_celery
+
+    for signal, handler in (
+        (before_task_publish, tb_celery.task_publish_handler),
+        (task_prerun, tb_celery.task_prerun_handler),
+        (task_postrun, tb_celery.task_postrun_handler),
+        (task_success, tb_celery.task_success_handler),
+        (task_failure, tb_celery.task_failure_handler),
+        (task_retry, tb_celery.task_retry_handler),
+    ):
+        signal.disconnect(handler)
     _assert_signals(check_is_connected=False)

tests/test_cli_list.py

@@ -0,0 +1,74 @@
+import json
+import os
+from http import HTTPStatus
+from unittest import mock
+
+import pytest
+from typer.testing import CliRunner
+
+from taskbadger.cli_main import app
+from taskbadger.internal.models import PaginatedTaskList
+from taskbadger.internal.types import Response
+from tests.utils import task_for_test
+
+runner = CliRunner()
+
+NEXT_URL = "https://taskbadger.net/api/org/project/tasks/?cursor=next-token"
+
+
+@pytest.fixture(autouse=True)
+def _mock_env():
+    with mock.patch.dict(
+        os.environ,
+        {
+            "TASKBADGER_ORG": "org",
+            "TASKBADGER_PROJECT": "project",
+            "TASKBADGER_API_KEY": "token",
+        },
+        clear=True,
+    ):
+        yield
+
+
+def _mock_list(*tasks, next_=None):
+    page = PaginatedTaskList(results=list(tasks), next_=next_, previous=None)
+    return Response(HTTPStatus.OK, b"", {}, page)
+
+
+def test_cli_list_pretty():
+    with mock.patch("taskbadger.sdk.task_list.sync_detailed") as list_:
+        # short id so the rich table doesn't truncate it at the default width
+        task = task_for_test(id="t1")
+        list_.return_value = _mock_list(task, next_=NEXT_URL)
+
+        result = runner.invoke(app, ["list"])
+
+    assert result.exit_code == 0, result.output
+    assert "t1" in result.output
+    assert task.name in result.output
+    assert "next-token" in result.output
+
+
+def test_cli_list_json():
+    with mock.patch("taskbadger.sdk.task_list.sync_detailed") as list_:
+        task = task_for_test()
+        list_.return_value = _mock_list(task, next_=NEXT_URL)
+
+        result = runner.invoke(app, ["list", "--format", "json"])
+
+    assert result.exit_code == 0, result.output
+    payload = json.loads(result.output)
+    assert payload["next_token"] == "next-token"
+    assert [t["id"] for t in payload["results"]] == [task.id]
+
+
+def test_cli_list_csv():
+    with mock.patch("taskbadger.sdk.task_list.sync_detailed") as list_:
+        task = task_for_test()
+        list_.return_value = _mock_list(task)
+
+        result = runner.invoke(app, ["list", "--format", "csv"])
+
+    assert result.exit_code == 0, result.output
+    assert task.id in result.output
+    assert "next_token" not in result.output

tests/test_parents.py

@@ -5,9 +5,11 @@
 is attached to the same root rather than to the child.
 """
 
+import copy
 import logging
 from unittest import mock
 
+import celery
 import procrastinate
 import pytest
 from procrastinate import testing
@@ -141,8 +143,15 @@ def test_list_tasks_filters_by_parent(httpx_mock):
         json={"next": None, "previous": None, "results": [_json_task_response(parent="parent_id")]},
         status_code=200,
     )
-    (child,) = list_tasks(parent="parent_id").results
+    tasks = list_tasks(parent="parent_id")
+    (child,) = tasks.results
+    assert isinstance(child, Task)
     assert child.parent == "parent_id"
+    assert list(tasks) == tasks.results
+    assert len(tasks) == 1
+    # a TaskList must survive copy / pickle: both probe for dunders that
+    # `__getattr__` must not try to delegate
+    assert len(copy.deepcopy(tasks)) == 1
 
 
 @pytest.mark.usefixtures("_bind_settings")
@@ -321,6 +330,106 @@ def test_celery_publish_explicit_parent_wins():
     assert create.call_args.kwargs["parent"] == "chosen"
 
 
+def _celery_app(**conf):
+    """A standalone app backed by in-memory transports, so no broker is needed."""
+    app = celery.Celery("test_parents", broker="memory://", backend="cache+memory://", **conf)
+
+    @app.task(bind=True, base=taskbadger.celery.Task, name="test_parents.add")
+    def add(self, a, b):
+        return a + b
+
+    return add
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_celery_apply_async_parent():
+    """`taskbadger_parent` on `apply_async` reaches the task created at publish time."""
+    add = _celery_app()
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.sdk.get_task"),
+    ):
+        create.return_value = task_for_test()
+        add.apply_async((2, 2), taskbadger_parent="chosen")
+
+    assert create.call_args.kwargs["parent"] == "chosen"
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_celery_apply_async_parent_beats_the_running_task():
+    add = _celery_app()
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.sdk.get_task"),
+    ):
+        create.return_value = task_for_test()
+        token = enter_task("root_id")
+        try:
+            add.apply_async((2, 2), taskbadger_parent="chosen")
+        finally:
+            exit_task(token)
+
+    assert create.call_args.kwargs["parent"] == "chosen"
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_celery_eager_apply_async_parent():
+    """Eager tasks are created in `task_prerun` rather than at publish time, but
+    the explicit parent still has to make it through."""
+    add = _celery_app(task_always_eager=True, task_eager_propagates=True)
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.celery.update_task_safe"),
+        mock.patch("taskbadger.sdk.get_task"),
+    ):
+        create.return_value = task_for_test()
+        assert add.apply_async((2, 2), taskbadger_parent="chosen").get() == 4
+
+    assert create.call_args.kwargs["parent"] == "chosen"
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_celery_eager_nests_under_the_running_task():
+    add = _celery_app(task_always_eager=True, task_eager_propagates=True)
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.celery.update_task_safe"),
+        mock.patch("taskbadger.sdk.get_task"),
+    ):
+        create.return_value = task_for_test()
+        token = enter_task("root_id")
+        try:
+            add.apply_async((2, 2))
+        finally:
+            exit_task(token)
+
+    assert create.call_args.kwargs["parent"] == "root_id"
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_celery_eager_explicit_none_parent_makes_a_root_task():
+    """`taskbadger_parent=None` asks for a root task, as it does at publish time."""
+    add = _celery_app(task_always_eager=True, task_eager_propagates=True)
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.celery.update_task_safe"),
+        mock.patch("taskbadger.sdk.get_task"),
+    ):
+        create.return_value = task_for_test()
+        token = enter_task("root_id")
+        try:
+            add.apply_async((2, 2), taskbadger_parent=None)
+        finally:
+            exit_task(token)
+
+    assert create.call_args.kwargs["parent"] is None
+
+
 # --- Procrastinate ------------------------------------------------------------