fix(routing): harden lifecycle coverage and callable task binding

- cover receiver and worker runtime cancellation and shutdown paths
- validate scheduler, router, registry, CLI, and broker lifecycle contracts
- support callable-object shared tasks during broker binding
- organize runtime fault tests into focused modules

Author: GefMar

taskiq/abc/broker.py

@@ -4,7 +4,7 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import AsyncGenerator, Awaitable, Callable
-from functools import wraps
+from functools import WRAPPER_ASSIGNMENTS, update_wrapper
 from logging import getLogger
 from typing import (
     TYPE_CHECKING,
@@ -13,6 +13,7 @@
     ParamSpec,
     TypeAlias,
     TypeVar,
+    cast,
     get_type_hints,
     overload,
 )
@@ -532,19 +533,29 @@ def _decorate_task(
         ):
             raise TypeError("base_cls must be a subclass of AsyncTaskiqDecoratedTask.")
 
-        sign = get_type_hints(metadata_source)
+        hint_source = metadata_source
+        if not hasattr(hint_source, "__annotations__"):
+            hint_source = cast(Any, metadata_source).__call__
+        sign = get_type_hints(hint_source)
         return_type = None
         if "return" in sign:
             return_type = sign["return"]
 
-        decorated_task = wraps(metadata_source)(
+        assigned_metadata = tuple(
+            attribute
+            for attribute in WRAPPER_ASSIGNMENTS
+            if hasattr(metadata_source, attribute)
+        )
+        decorated_task = update_wrapper(
             task_cls(
                 broker=self,
                 original_func=func,
                 labels=labels,
                 task_name=real_task_name,
                 return_type=return_type,  # type: ignore
             ),
+            metadata_source,
+            assigned=assigned_metadata,
         )
 
         if register:

tests/abc/test_broker.py

@@ -68,6 +68,18 @@ async def test_func() -> None:
     }
 
 
+def test_decorator_generates_unique_lambda_task_names() -> None:
+    """Anonymous tasks must not collide in the broker task registry."""
+    broker = _TestBroker()
+
+    first_task = broker.task(lambda: None)
+    second_task = broker.task(lambda: None)
+
+    assert ":lambda_" in first_task.task_name
+    assert ":lambda_" in second_task.task_name
+    assert first_task.task_name != second_task.task_name
+
+
 def test_kicker_labels_modification() -> None:
     """Test that using kicker.with_labels doesn't modify task's labels globally."""
     broker = _TestBroker()

tests/brokers/test_inmemory.py

@@ -1,10 +1,11 @@
 import asyncio
+import logging
 import uuid
 from typing import Any
 
 import pytest
 
-from taskiq import Flow, InMemoryBroker
+from taskiq import Flow, InMemoryBroker, TaskiqMessage
 from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.brokers.inmemory_broker import InmemoryResultBackend
 from taskiq.events import TaskiqEvents
@@ -39,16 +40,18 @@ async def shutdown(self) -> None:
 
 
 class RecordingResultBackend(InmemoryResultBackend[Any]):
-    """Record result backend lifecycle calls and optional startup failure."""
+    """Record result backend lifecycle calls and optional failures."""
 
     def __init__(
         self,
         events: list[str],
         startup_error: Exception | None = None,
+        shutdown_error: Exception | None = None,
     ) -> None:
         super().__init__()
         self.events = events
         self.startup_error = startup_error
+        self.shutdown_error = shutdown_error
 
     async def startup(self) -> None:
         """Record startup and raise the configured error."""
@@ -59,6 +62,20 @@ async def startup(self) -> None:
     async def shutdown(self) -> None:
         """Record shutdown."""
         self.events.append("backend.shutdown")
+        if self.shutdown_error is not None:
+            raise self.shutdown_error
+
+
+class FailingExecutionMiddleware(TaskiqMiddleware):
+    """Raise before execution to expose background callback failures."""
+
+    def __init__(self, error: Exception) -> None:
+        super().__init__()
+        self.error = error
+
+    def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage:
+        """Raise the configured callback failure."""
+        raise self.error
 
 
 async def test_inmemory_success() -> None:
@@ -197,6 +214,48 @@ async def test_shutdown_closes_remaining_resources_after_middleware_failure() ->
         broker.executor.submit(int)
 
 
+async def test_shutdown_keeps_first_failure_and_closes_every_resource(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    events: list[str] = []
+    event_error = LifecycleError("event shutdown failed")
+    middleware_error = LifecycleError("middleware shutdown failed")
+    backend_error = LifecycleError("backend shutdown failed")
+    broker = InMemoryBroker()
+    broker.with_middlewares(
+        RecordingLifecycleMiddleware(events, shutdown_error=middleware_error),
+    )
+    broker.with_result_backend(
+        RecordingResultBackend(events, shutdown_error=backend_error),
+    )
+
+    @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN)
+    def fail_client_shutdown(state: TaskiqState) -> None:
+        events.append("client.shutdown")
+        raise event_error
+
+    @broker.on_event(TaskiqEvents.WORKER_SHUTDOWN)
+    def record_worker_shutdown(state: TaskiqState) -> None:
+        events.append("worker.shutdown")
+
+    caplog.set_level(logging.ERROR, logger="taskiq")
+    await broker.startup()
+
+    with pytest.raises(LifecycleError) as exc_info:
+        await broker.shutdown()
+
+    assert exc_info.value is event_error
+    assert events == [
+        "middleware.startup",
+        "backend.startup",
+        "client.shutdown",
+        "worker.shutdown",
+        "middleware.shutdown",
+        "backend.shutdown",
+    ]
+    assert caplog.text.count("Additional error while shutting down") == 2
+
+
 async def test_shutdown() -> None:
     broker = InMemoryBroker()
     test_value = uuid.uuid4().hex
@@ -265,6 +324,54 @@ async def test_task() -> None:
     assert not broker._running_tasks
 
 
+async def test_wait_all_propagates_task_failure() -> None:
+    task_error = LifecycleError("in-memory task failed")
+    broker = InMemoryBroker()
+    broker.with_middlewares(FailingExecutionMiddleware(task_error))
+
+    @broker.task
+    async def failing_task() -> None:
+        return None
+
+    await failing_task.kiq()
+
+    with pytest.raises(LifecycleError) as exc_info:
+        await broker.wait_all()
+
+    assert exc_info.value is task_error
+    assert not broker._running_tasks
+
+
+async def test_shutdown_cleans_resources_after_task_failure() -> None:
+    events: list[str] = []
+    task_error = LifecycleError("in-memory task failed")
+    broker = InMemoryBroker()
+    broker.with_middlewares(RecordingLifecycleMiddleware(events))
+    broker.with_middlewares(FailingExecutionMiddleware(task_error))
+    broker.with_result_backend(RecordingResultBackend(events))
+
+    @broker.task
+    async def failing_task() -> None:
+        return None
+
+    await broker.startup()
+    await failing_task.kiq()
+
+    with pytest.raises(LifecycleError) as exc_info:
+        await broker.shutdown()
+
+    assert exc_info.value is task_error
+    assert events == [
+        "middleware.startup",
+        "backend.startup",
+        "middleware.shutdown",
+        "backend.shutdown",
+    ]
+    assert not broker._running_tasks
+    with pytest.raises(RuntimeError, match="cannot schedule new futures"):
+        broker.executor.submit(int)
+
+
 async def test_shutdown_waits_for_running_tasks_before_resource_cleanup() -> None:
     task_started = asyncio.Event()
     finish_task = asyncio.Event()

tests/cli/worker/test_broker_selection.py

@@ -1,10 +1,15 @@
+import asyncio
 import logging
+from concurrent.futures import Executor
 
 import pytest
 
 from taskiq import InMemoryBroker, TaskiqRouter
+from taskiq.acks import AcknowledgeType
 from taskiq.cli.worker.args import WorkerArgs
-from taskiq.cli.worker.run import get_broker_selection, shutdown_broker
+from taskiq.cli.worker.run import get_broker_selection, shutdown_broker, start_listen
+from taskiq.receiver import Receiver
+from taskiq.receiver.runtime import BrokerSelection, WorkerRuntimeOptions
 from taskiq.warnings import TaskiqDeprecationWarning
 
 
@@ -17,6 +22,103 @@ def test_get_broker_selection_preserves_direct_single_broker() -> None:
     assert not selection.explicit_group
 
 
+def test_start_listen_builds_runtime_from_worker_arguments(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    broker = InMemoryBroker()
+    args = WorkerArgs(
+        broker=broker,
+        modules=["application.tasks"],
+        tasks_pattern=("**/jobs.py",),
+        fs_discover=True,
+        configure_logging=False,
+        max_threadpool_threads=2,
+        no_parse=True,
+        shutdown_timeout=0.7,
+        max_async_tasks=7,
+        max_async_tasks_jitter=2,
+        receiver_arg=[("dependency", "value")],
+        max_prefetch=3,
+        no_propagate_errors=True,
+        ack_type=AcknowledgeType.WHEN_RECEIVED,
+        max_tasks_per_child=11,
+        wait_tasks_timeout=0.5,
+    )
+    loop = asyncio.new_event_loop()
+    captured_selection: BrokerSelection | None = None
+    captured_receiver_type: type[Receiver] | None = None
+    captured_executor: Executor | None = None
+    captured_options: WorkerRuntimeOptions | None = None
+    runtime_ran = False
+    import_calls: list[tuple[list[str], tuple[str, ...], bool]] = []
+
+    class RecordingRuntime:
+        def __init__(
+            self,
+            *,
+            selection: BrokerSelection,
+            receiver_type: type[Receiver],
+            executor: Executor,
+            options: WorkerRuntimeOptions,
+        ) -> None:
+            nonlocal captured_selection
+            nonlocal captured_receiver_type
+            nonlocal captured_executor
+            nonlocal captured_options
+            captured_selection = selection
+            captured_receiver_type = receiver_type
+            captured_executor = executor
+            captured_options = options
+
+        async def run(self, finish_event: asyncio.Event) -> None:
+            nonlocal runtime_ran
+            runtime_ran = True
+            finish_event.set()
+
+    def record_import_tasks(
+        modules: list[str],
+        task_pattern: tuple[str, ...],
+        fs_discover: bool,
+    ) -> None:
+        import_calls.append((modules, task_pattern, fs_discover))
+
+    monkeypatch.setattr("taskiq.cli.worker.run.uvloop", None)
+    monkeypatch.setattr("taskiq.cli.worker.run.signal.signal", lambda *args: None)
+    monkeypatch.setattr("taskiq.cli.worker.run.asyncio.new_event_loop", lambda: loop)
+    monkeypatch.setattr("taskiq.cli.worker.run.import_tasks", record_import_tasks)
+    monkeypatch.setattr(
+        "taskiq.cli.worker.run.get_receiver_type",
+        lambda args: Receiver,
+    )
+    monkeypatch.setattr("taskiq.cli.worker.run.WorkerRuntime", RecordingRuntime)
+
+    try:
+        start_listen(args)
+    finally:
+        asyncio.set_event_loop(None)
+        loop.close()
+
+    assert captured_selection is not None
+    assert captured_selection.listeners == (broker,)
+    assert broker.is_worker_process
+    assert captured_receiver_type is Receiver
+    assert captured_executor is not None
+    assert captured_options == WorkerRuntimeOptions(
+        validate_params=False,
+        max_async_tasks=7,
+        max_async_tasks_jitter=2,
+        max_prefetch=3,
+        propagate_exceptions=False,
+        ack_type=AcknowledgeType.WHEN_RECEIVED,
+        max_tasks_to_execute=11,
+        wait_tasks_timeout=0.5,
+        shutdown_timeout=0.7,
+        receiver_kwargs={"dependency": "value"},
+    )
+    assert import_calls == [(["application.tasks"], ("**/jobs.py",), True)]
+    assert runtime_ran
+
+
 @pytest.mark.parametrize("use_factory", [False, True])
 def test_get_broker_selection_preserves_imported_single_broker(
     monkeypatch: pytest.MonkeyPatch,
@@ -67,6 +169,25 @@ async def shutdown(self) -> None:
     assert diagnostic not in caplog.text
 
 
+async def test_legacy_shutdown_broker_helper_preserves_cancellation() -> None:
+    shutdown_error = asyncio.CancelledError("adapter shutdown cancelled")
+
+    class CancelledShutdownBroker(InMemoryBroker):
+        async def shutdown(self) -> None:
+            raise shutdown_error
+
+    broker = CancelledShutdownBroker()
+
+    with (
+        pytest.warns(
+            TaskiqDeprecationWarning,
+            match="WorkerRuntime",
+        ),
+        pytest.raises(asyncio.CancelledError),
+    ):
+        await shutdown_broker(broker, timeout=1)
+
+
 @pytest.mark.parametrize("use_factory", [False, True])
 def test_get_broker_selection_imports_explicit_listener_sequence(
     monkeypatch: pytest.MonkeyPatch,

tests/receiver/conftest.py

@@ -0,0 +1,11 @@
+from collections.abc import Iterator
+from concurrent.futures import ThreadPoolExecutor
+
+import pytest
+
+
+@pytest.fixture
+def runtime_executor() -> Iterator[ThreadPoolExecutor]:
+    """Provide the process-wide sync executor owned by a worker runtime."""
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        yield executor