diff --git a/taskiq/abc/broker.py b/taskiq/abc/broker.py index 41813f46..8e47291c 100644 --- a/taskiq/abc/broker.py +++ b/taskiq/abc/broker.py @@ -167,7 +167,7 @@ async def listen( def task( self, task_name: Callable[_FuncParams, _ReturnType], - ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: + ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: # pragma: no cover ... @overload @@ -178,7 +178,7 @@ def task( ) -> Callable[ [Callable[_FuncParams, _ReturnType]], AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType], - ]: + ]: # pragma: no cover ... def task( # type: ignore[misc] @@ -222,7 +222,7 @@ def inner( nonlocal inner_task_name # noqa: WPS420 if inner_task_name is None: fmodule = func.__module__ - if fmodule == "__main__": + if fmodule == "__main__": # pragma: no cover fmodule = ".".join( # noqa: WPS220 sys.argv[0] .removesuffix( diff --git a/taskiq/abc/tests/test_broker.py b/taskiq/abc/tests/test_broker.py new file mode 100644 index 00000000..631a08c1 --- /dev/null +++ b/taskiq/abc/tests/test_broker.py @@ -0,0 +1,66 @@ +from typing import Any, Callable, Coroutine + +from taskiq.abc.broker import AsyncBroker +from taskiq.decor import AsyncTaskiqDecoratedTask +from taskiq.message import BrokerMessage + + +class _TestBroker(AsyncBroker): + """Broker for testing purpose.""" + + async def kick(self, message: BrokerMessage) -> None: + """ + This method is used to send messages. + + But in this case it just throws messages away. + + :param message: message to lost. + """ + + async def listen( + self, + callback: Callable[[BrokerMessage], Coroutine[Any, Any, None]], + ) -> None: + """ + This method is not implemented. + + :param callback: callback that is never called. + """ + + +def test_decorator_success() -> None: + """Test that decoration without parameters works.""" + tbrok = _TestBroker() + + @tbrok.task + async def test_func() -> None: + """Some test function.""" + + assert isinstance(test_func, AsyncTaskiqDecoratedTask) + + +def test_decorator_with_name_success() -> None: + """Test that task_name is successfully set.""" + tbrok = _TestBroker() + + @tbrok.task(task_name="my_task") + async def test_func() -> None: + """Some test function.""" + + assert isinstance(test_func, AsyncTaskiqDecoratedTask) + assert test_func.task_name == "my_task" + + +def test_decorator_with_labels_success() -> None: + """Tests that labels are assigned for task as is.""" + tbrok = _TestBroker() + + @tbrok.task(label1=1, label2=2) + async def test_func() -> None: + """Some test function.""" + + assert isinstance(test_func, AsyncTaskiqDecoratedTask) + assert test_func.labels == { + "label1": 1, + "label2": 2, + }