fix: sync decorators (#343)

Co-authored-by: Anton <myprojectorterrypratchett@gmail.com>
Co-authored-by: Dmitrii Anfimov <anfimov@tochka.com>

Author: 3 people

taskiq/receiver/receiver.py

@@ -251,8 +251,18 @@ async def run_task(  # noqa: C901, PLR0912, PLR0915
             if timeout is not None:
                 if not is_coroutine:
                     logger.warning("Timeouts for sync tasks don't work in python well.")
-                target_future = asyncio.wait_for(target_future, float(timeout))
-            returned = await target_future
+
+                with anyio.fail_after(float(timeout)):
+                    target_future = await target_future
+                    if inspect.isawaitable(target_future):
+                        target_future = await target_future
+
+            else:
+                target_future = await target_future
+                if inspect.isawaitable(target_future):
+                    target_future = await target_future
+
+            returned = target_future
         except NoResultError as no_res_exc:
             found_exception = no_res_exc
             logger.warning(

tests/receiver/test_receiver.py

@@ -4,6 +4,7 @@
 import time
 from collections.abc import Generator
 from concurrent.futures import ThreadPoolExecutor
+from functools import wraps
 from typing import Any, ClassVar
 
 import pytest
@@ -516,3 +517,30 @@ async def test_func() -> int:
         ),
     )
     assert result.return_value == EXPECTED_CTX_VALUE
+
+
+async def test_sync_decorator_on_async_function() -> None:
+    broker = InMemoryBroker()
+    wrapper_call = False
+
+    def wrapper(f: Any) -> Any:
+        @wraps(f)
+        def wrapper_impl(*args: Any, **kwargs: Any) -> Any:
+            nonlocal wrapper_call
+
+            wrapper_call = True
+            return f(*args, **kwargs)
+
+        return wrapper_impl
+
+    @broker.task
+    @wrapper
+    async def task_no_result() -> str:
+        return "some value"
+
+    task = await task_no_result.kiq()
+    resp = await task.wait_result(timeout=1)
+
+    assert resp.return_value == "some value"
+    assert not broker._running_tasks
+    assert wrapper_call is True