test_task.py

import logging
import uuid
from typing import Any, TypeVar

import pytest
from pydantic import BaseModel

from taskiq import serializers
from taskiq.abc import AsyncResultBackend
from taskiq.abc.serializer import TaskiqSerializer
from taskiq.result import TaskiqResult
from taskiq.task import AsyncTaskiqTask

_ReturnType = TypeVar("_ReturnType")


class SerializingBackend(AsyncResultBackend[_ReturnType]):
    def __init__(self, serializer: TaskiqSerializer) -> None:
        self._serializer = serializer
        self._results: dict[str, bytes] = {}

    async def set_result(
        self,
        task_id: str,
        result: TaskiqResult[_ReturnType],  # type: ignore
    ) -> None:
        """Set result with dumping."""
        self._results[task_id] = self._serializer.dumpb(result.model_dump(mode="json"))

    async def is_result_ready(self, task_id: str) -> bool:
        """Check if result is ready."""
        return task_id in self._results

    async def get_result(
        self,
        task_id: str,
        with_logs: bool = False,
    ) -> TaskiqResult[_ReturnType]:
        """Get result with loading."""
        data = self._results[task_id]
        return TaskiqResult.model_validate(self._serializer.loadb(data))


@pytest.mark.parametrize(
    "serializer",
    [
        serializers.MSGPackSerializer(),
        serializers.CBORSerializer(),
        serializers.PickleSerializer(),
        serializers.JSONSerializer(),
    ],
)
async def test_res_parsing_success(serializer: TaskiqSerializer) -> None:
    class MyResult(BaseModel):
        name: str
        age: int

    res = MyResult(name="test", age=10)
    res_back: AsyncResultBackend[MyResult] = SerializingBackend(serializer)
    test_id = str(uuid.uuid4())
    await res_back.set_result(
        test_id,
        TaskiqResult(
            is_err=False,
            return_value=res,
            execution_time=0.0,
        ),
    )
    sent_task = AsyncTaskiqTask(test_id, res_back, MyResult)
    parsed = await sent_task.wait_result()
    assert isinstance(parsed.return_value, MyResult)


async def test_res_parsing_skipped_on_error(caplog: pytest.LogCaptureFixture) -> None:
    res_back: AsyncResultBackend[Any] = SerializingBackend(serializers.JSONSerializer())
    test_id = str(uuid.uuid4())
    await res_back.set_result(
        test_id,
        TaskiqResult(
            is_err=True,
            return_value=None,
            execution_time=0.0,
        ),
    )
    sent_task: AsyncTaskiqTask[Any] = AsyncTaskiqTask(test_id, res_back, int)
    with caplog.at_level(logging.WARNING, logger="taskiq.task"):
        parsed = await sent_task.wait_result()
    assert parsed.is_err
    assert parsed.return_value is None
    assert "Cannot parse return type" not in caplog.text


async def test_res_parsing_failure(caplog: pytest.LogCaptureFixture) -> None:
    res_back: AsyncResultBackend[Any] = SerializingBackend(serializers.JSONSerializer())
    test_id = str(uuid.uuid4())
    await res_back.set_result(
        test_id,
        TaskiqResult(
            is_err=False,
            return_value="not-an-int",
            execution_time=0.0,
        ),
    )
    sent_task: AsyncTaskiqTask[Any] = AsyncTaskiqTask(test_id, res_back, int)
    with caplog.at_level(logging.WARNING, logger="taskiq.task"):
        parsed = await sent_task.wait_result()
    assert not parsed.is_err
    assert parsed.return_value == "not-an-int"
    assert "Cannot parse return type" in caplog.text