fix: TaskiqAdminMiddleware work with dataclasses

Author: danfimov

pyproject.toml

@@ -68,6 +68,7 @@ dev = [
 	"freezegun>=1.5.5",
 	"tzdata>=2025.2; sys_platform == 'win32'",
 	"opentelemetry-test-utils (>=0.59b0,<1)",
+	"polyfactory>=3.1.0",
 ]
 
 [project.urls]
@@ -172,8 +173,8 @@ lint.ignore = [
 	"PLR0913", # Too many arguments for function call
 	"D106",    # Missing docstring in public nested class
 ]
-exclude = [".venv/"]
 lint.mccabe = { max-complexity = 10 }
+exclude = [".venv/"]
 line-length = 88
 
 [tool.ruff.lint.per-file-ignores]

taskiq/middlewares/taskiq_admin_middleware.py

@@ -7,6 +7,7 @@
 import aiohttp
 
 from taskiq.abc.middleware import TaskiqMiddleware
+from taskiq.compat import model_dump
 from taskiq.message import TaskiqMessage
 from taskiq.result import TaskiqResult
 
@@ -115,12 +116,13 @@ async def post_send(self, message: TaskiqMessage) -> None:
 
         :param message: kicked message.
         """
+        dict_message: dict[str, Any] = model_dump(message)
         await self._spawn_request(
             f"/api/tasks/{message.task_id}/queued",
             {
-                "args": message.args,
-                "kwargs": message.kwargs,
-                "labels": message.labels,
+                "args": dict_message["args"],
+                "kwargs": dict_message["kwargs"],
+                "labels": dict_message["labels"],
                 "queuedAt": self._now_iso(),
                 "taskName": message.task_name,
                 "worker": self.__ta_broker_name,
@@ -137,12 +139,13 @@ async def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage:
         :param message: incoming parsed taskiq message.
         :return: modified message.
         """
+        dict_message: dict[str, Any] = model_dump(message)
         await self._spawn_request(
             f"/api/tasks/{message.task_id}/started",
             {
-                "args": message.args,
-                "kwargs": message.kwargs,
-                "labels": message.labels,
+                "args": dict_message["args"],
+                "kwargs": dict_message["kwargs"],
+                "labels": dict_message["labels"],
                 "startedAt": self._now_iso(),
                 "taskName": message.task_name,
                 "worker": self.__ta_broker_name,
@@ -164,12 +167,13 @@ async def post_execute(
         :param message: incoming message.
         :param result: result of execution for current task.
         """
+        dict_result: dict[str, Any] = model_dump(result)
         await self._spawn_request(
             f"/api/tasks/{message.task_id}/executed",
             {
                 "finishedAt": self._now_iso(),
                 "executionTime": result.execution_time,
                 "error": None if result.error is None else repr(result.error),
-                "returnValue": {"return_value": result.return_value},
+                "returnValue": {"return_value": dict_result["return_value"]},
             },
         )

tests/middlewares/admin_middleware/conftest.py

@@ -0,0 +1,72 @@
+import pytest
+from aiohttp import web
+from aiohttp.test_utils import TestServer
+from typing_extensions import AsyncGenerator
+
+from taskiq.brokers.inmemory_broker import InMemoryBroker
+from taskiq.brokers.shared_broker import async_shared_broker
+from taskiq.middlewares import TaskiqAdminMiddleware
+from tests.middlewares.admin_middleware.dto import (
+    DataclassDTO,
+    PydanticDTO,
+    TypedDictDTO,
+)
+
+
+@pytest.fixture(scope="session")
+async def admin_api_server() -> AsyncGenerator[TestServer, None]:
+    async def handle_queued(request: web.Request) -> web.Response:
+        return web.json_response({"status": "ok"}, status=200)
+
+    async def handle_started(request: web.Request) -> web.Response:
+        return web.json_response({"status": "ok"}, status=200)
+
+    async def handle_executed(request: web.Request) -> web.Response:
+        return web.json_response({"status": "ok"}, status=200)
+
+    app = web.Application()
+    app.router.add_post("/api/tasks/{task_id}/queued", handle_queued)
+    app.router.add_post("/api/tasks/{task_id}/started", handle_started)
+    app.router.add_post("/api/tasks/{task_id}/executed", handle_executed)
+
+    server = TestServer(app)
+    await server.start_server()
+
+    yield server
+
+    # Останавливаем сервер после теста
+    await server.close()
+
+
+@pytest.fixture
+async def broker_with_admin_middleware(
+    admin_api_server: TestServer,
+) -> AsyncGenerator[InMemoryBroker, None]:
+    broker = InMemoryBroker().with_middlewares(
+        TaskiqAdminMiddleware(
+            str(admin_api_server.make_url("/")),  # URL тестового сервера
+            "supersecret",
+            taskiq_broker_name="InMemory",
+        ),
+    )
+
+    broker.register_task(task_with_dataclass, task_name="task_with_dataclass")
+    broker.register_task(task_with_typed_dict, task_name="task_with_typed_dict")
+    broker.register_task(task_with_pydantic_model, task_name="task_with_pydantic_model")
+    async_shared_broker.default_broker(broker)
+
+    await broker.startup()
+    yield broker
+    await broker.shutdown()
+
+
+async def task_with_dataclass(dto: DataclassDTO) -> None:
+    assert dto
+
+
+async def task_with_typed_dict(dto: TypedDictDTO) -> None:
+    assert dto
+
+
+async def task_with_pydantic_model(dto: PydanticDTO) -> None:
+    assert dto

tests/middlewares/admin_middleware/dto.py

@@ -0,0 +1,36 @@
+from dataclasses import dataclass
+from typing import TypedDict
+
+import pydantic
+
+
+@dataclass(frozen=True, slots=True)
+class DataclassNestedDTO:
+    id: int
+    name: str
+
+
+@dataclass(frozen=True, slots=True)
+class DataclassDTO:
+    nested: DataclassNestedDTO
+    recipients: list[str]
+    subject: str
+    attachments: list[str] | None = None
+    text: str | None = None
+    html: str | None = None
+
+
+class PydanticDTO(pydantic.BaseModel):
+    number: int
+    text: str
+    flag: bool
+    list: list[float]
+    dictionary: dict[str, str] | None = None
+
+
+class TypedDictDTO(TypedDict):
+    id: int
+    name: str
+    active: bool
+    scores: list[int]
+    metadata: dict[str, str] | None

tests/middlewares/admin_middleware/test_arguments_formatting.py

@@ -0,0 +1,52 @@
+import pytest
+from polyfactory.factories import BaseFactory, DataclassFactory, TypedDictFactory
+from polyfactory.factories.pydantic_factory import ModelFactory
+
+from taskiq.brokers.inmemory_broker import InMemoryBroker
+from tests.middlewares.admin_middleware.dto import (
+    DataclassDTO,
+    PydanticDTO,
+    TypedDictDTO,
+)
+
+
+class DataclassDTOFactory(DataclassFactory[DataclassDTO]):
+    __model__ = DataclassDTO
+
+
+class TypedDictDTOFactory(TypedDictFactory[TypedDictDTO]):
+    __model__ = TypedDictDTO
+
+
+class PydanticDTOFactory(ModelFactory[PydanticDTO]):
+    __model__ = PydanticDTO
+
+
+class TestArgumentsFormattingInAdminMiddleware:
+    @pytest.mark.parametrize(
+        "dto_factory, task_name",
+        [
+            pytest.param(DataclassDTOFactory, "task_with_dataclass", id="dataclass"),
+            pytest.param(TypedDictDTOFactory, "task_with_typed_dict", id="typeddict"),
+            pytest.param(PydanticDTOFactory, "task_with_pydantic_model", id="pydantic"),
+        ],
+    )
+    async def test_when_task_dto_passed__then_middleware_successfully_send_request(
+        self,
+        broker_with_admin_middleware: InMemoryBroker,
+        dto_factory: type[BaseFactory],  # type: ignore[type-arg]
+        task_name: str,
+    ) -> None:
+        # given
+        task_arguments = dto_factory.build()
+        task = broker_with_admin_middleware.find_task(task_name)
+        assert task is not None, f"Task {task_name} should be registered in the broker"
+
+        # when
+        kicked_task = await task.kiq(task_arguments)
+        await broker_with_admin_middleware.wait_all()
+
+        # then
+        result = await kicked_task.get_result()
+        # we just expect no errors during post_send/pre_execute/post_execute
+        assert result.error is None