Merge pull request #5 from Lash-L/asyncio_timeout

Change to timeout from wait_for

Author: humbertogontijo

pyproject.toml

@@ -14,6 +14,7 @@ roborock = "roborock.cli:main"
 python = "^3.8"
 click = ">=8"
 aiohttp = "*"
+async-timeout = "*"
 pycryptodome = "~3.16.0"
 pycryptodomex = {version = "~3.16.0", markers = "sys_platform == 'darwin'"}
 paho-mqtt = "~1.6.1"

roborock/roborock_queue.py

@@ -1,6 +1,7 @@
 import asyncio
 from asyncio import Queue
 from typing import Any
+import async_timeout
 
 from .exceptions import RoborockException
 
@@ -12,7 +13,9 @@ def __init__(self, protocol: int, *args):
         self.protocol = protocol
 
     async def async_put(self, item: tuple[Any, RoborockException | None], timeout: float | int) -> None:
-        return await asyncio.wait_for(self.put(item), timeout=timeout)
+        async with async_timeout.timeout(timeout):
+            await self.put(item)
 
     async def async_get(self, timeout: float | int) -> tuple[Any, RoborockException | None]:
-        return await asyncio.wait_for(self.get(), timeout=timeout)
+        async with async_timeout.timeout(timeout):
+            return await self.get()