taskiq/tests/cli/worker/test_parameters_parsing.py at 8f97d38c0f4568292fbc039f409ad01664640343 · taskiq-python/taskiq · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
import inspect
from dataclasses import dataclass
from typing import Any, Type, get_type_hints

import pytest
from pydantic import BaseModel

from taskiq.compat import model_copy
from taskiq.message import TaskiqMessage
from taskiq.receiver.params_parser import parse_params


class _TestPydanticClass(BaseModel):
    field: str


@dataclass
class _TestDataclass:
    field: str


def test_parse_params_no_signature() -> None:
    """Test that params aren't parsed if no annotation is supplied."""
    src_msg = TaskiqMessage(
        task_id="",
        task_name="",
        labels={},
        queue="taksiq",
        args=[1, 2],
        kwargs={"a": 1},
    )
    modify_msg = model_copy(src_msg, deep=True)
    parse_params(
        signature=None,
        type_hints={},
        message=modify_msg,
    )

    assert modify_msg == src_msg


@pytest.mark.parametrize("test_class", [_TestPydanticClass, _TestDataclass])
def test_parse_params_classes(test_class: Type[Any]) -> None:
    """Test that dataclasses are parsed correctly."""

    def test_func(param: test_class) -> test_class:  # type: ignore
        return param

    msg_with_args = TaskiqMessage(
        task_id="",
        task_name="",
        labels={},
        queue="taksiq",
        args=[{"field": "test_val"}],
        kwargs={},
    )

    parse_params(
        inspect.signature(test_func),
        get_type_hints(test_func),
        msg_with_args,
    )

    assert isinstance(msg_with_args.args[0], test_class)
    assert msg_with_args.args[0].field == "test_val"

    msg_with_kwargs = TaskiqMessage(
        task_id="",
        task_name="",
        labels={},
        queue="taksiq",
        args=[],
        kwargs={"param": {"field": "test_val"}},
    )

    parse_params(
        inspect.signature(test_func),
        get_type_hints(test_func),
        msg_with_kwargs,
    )

    assert isinstance(msg_with_kwargs.kwargs["param"], test_class)
    assert msg_with_kwargs.kwargs["param"].field == "test_val"


@pytest.mark.parametrize("test_class", [_TestPydanticClass, _TestDataclass])
def test_parse_params_wrong_data(test_class: Type[Any]) -> None:
    """Tests that wrong data isn't parsed and doesn't throw errors."""

    def test_func(param: test_class) -> test_class:  # type: ignore
        return param

    msg_with_args = TaskiqMessage(
        task_id="",
        task_name="",
        labels={},
        queue="taksiq",
        args=[{"unknown": "unknown"}],
        kwargs={},
    )

    parse_params(
        inspect.signature(test_func),
        get_type_hints(test_func),
        msg_with_args,
    )

    assert isinstance(msg_with_args.args[0], dict)

    msg_with_kwargs = TaskiqMessage(
        task_id="",
        task_name="",
        labels={},
        queue="taksiq",
        args=[],
        kwargs={"param": {"unknown": "unknown"}},
    )

    parse_params(
        inspect.signature(test_func),
        get_type_hints(test_func),
        msg_with_kwargs,
    )

    assert isinstance(msg_with_kwargs.kwargs["param"], dict)


@pytest.mark.parametrize("test_class", [_TestPydanticClass, _TestDataclass])
def test_parse_params_nones(test_class: Type[Any]) -> None:
    """Tests that None values are not parsed."""

    def test_func(param: test_class) -> test_class:  # type: ignore
        return param

    msg_with_args = TaskiqMessage(
        task_id="",
        task_name="",
        queue="taksiq",
        labels={},
        args=[None],
        kwargs={},
    )

    parse_params(inspect.signature(test_func), get_type_hints(test_func), msg_with_args)

    assert msg_with_args.args[0] is None

    msg_with_kwargs = TaskiqMessage(
        task_id="",
        task_name="",
        queue="taksiq",
        labels={},
        args=[],
        kwargs={"param": None},
    )

    parse_params(
        inspect.signature(test_func),
        get_type_hints(test_func),
        msg_with_kwargs,
    )

    assert msg_with_kwargs.kwargs["param"] is None