Replace all asyncio.wait_fors with asyncio.timeout (#546)

bellows/ezsp/__init__.py

@@ -5,9 +5,15 @@
 import asyncio
 import functools
 import logging
+import sys
 from typing import Any, Callable, Dict, List, Tuple, Union
 import urllib.parse
 
+if sys.version_info[:2] < (3, 11):
+    from async_timeout import timeout as asyncio_timeout  # pragma: no cover
+else:
+    from asyncio import timeout as asyncio_timeout  # pragma: no cover
+
 import zigpy.config
 
 import bellows.config as conf
@@ -58,7 +64,9 @@ async def probe(cls, device_config: Dict) -> bool | dict[str, int | str | bool]:
         ):
             ezsp = cls(conf.SCHEMA_DEVICE(config))
             try:
-                await asyncio.wait_for(ezsp._probe(), timeout=PROBE_TIMEOUT)
+                async with asyncio_timeout(PROBE_TIMEOUT):
+                    await ezsp._probe()
+
                 return config
             except Exception as exc:
                 LOGGER.debug(
@@ -83,10 +91,8 @@ async def _startup_reset(self):
         parsed_path = urllib.parse.urlparse(self._config[conf.CONF_DEVICE_PATH])
         if parsed_path.scheme == "socket":
             try:
-                await asyncio.wait_for(
-                    self._gw.wait_for_startup_reset(),
-                    NETWORK_COORDINATOR_STARTUP_RESET_WAIT,
-                )
+                async with asyncio_timeout(NETWORK_COORDINATOR_STARTUP_RESET_WAIT):
+                    await self._gw.wait_for_startup_reset()
             except asyncio.TimeoutError:
                 pass
             else:
@@ -233,8 +239,9 @@ def cb(frame_name: str, response: List) -> None:
             (status,) = await self._command("leaveNetwork")
             if status != t.EmberStatus.SUCCESS:
                 raise EzspError(f"failed to leave network: {status.name}")
-            result = await asyncio.wait_for(stack_status, timeout=timeout)
-            return result
+
+            async with asyncio_timeout(timeout):
+                return await stack_status
         finally:
             self.remove_callback(cb_id)
 

bellows/ezsp/protocol.py

@@ -3,8 +3,14 @@
 import binascii
 import functools
 import logging
+import sys
 from typing import Any, Callable, Dict, Optional, Tuple
 
+if sys.version_info[:2] < (3, 11):
+    from async_timeout import timeout as asyncio_timeout  # pragma: no cover
+else:
+    from asyncio import timeout as asyncio_timeout  # pragma: no cover
+
 from bellows.config import CONF_EZSP_CONFIG, CONF_EZSP_POLICIES
 from bellows.exception import EzspError
 from bellows.typing import GatewayType
@@ -111,7 +117,7 @@ async def get_free_buffers(self) -> Optional[int]:
 
         return int.from_bytes(value, byteorder="little")
 
-    def command(self, name, *args) -> asyncio.Future:
+    async def command(self, name, *args) -> Any:
         """Serialize command and send it."""
         LOGGER.debug("Send command %s: %s", name, args)
         data = self._ezsp_frame(name, *args)
@@ -120,7 +126,9 @@ def command(self, name, *args) -> asyncio.Future:
         future = asyncio.Future()
         self._awaiting[self._seq] = (c[0], c[2], future)
         self._seq = (self._seq + 1) % 256
-        return asyncio.wait_for(future, timeout=EZSP_CMD_TIMEOUT)
+
+        async with asyncio_timeout(EZSP_CMD_TIMEOUT):
+            return await future
 
     async def set_source_routing(self) -> None:
         """Enable source routing on NCP."""

bellows/uart.py

@@ -1,6 +1,12 @@
 import asyncio
 import binascii
 import logging
+import sys
+
+if sys.version_info[:2] < (3, 11):
+    from async_timeout import timeout as asyncio_timeout  # pragma: no cover
+else:
+    from asyncio import timeout as asyncio_timeout  # pragma: no cover
 
 import zigpy.serial
 
@@ -257,7 +263,9 @@ async def reset(self):
         self._reset_future = asyncio.get_event_loop().create_future()
         self._reset_future.add_done_callback(self._reset_cleanup)
         self.write(self._rst_frame())
-        return await asyncio.wait_for(self._reset_future, timeout=RESET_TIMEOUT)
+
+        async with asyncio_timeout(RESET_TIMEOUT):
+            return await self._reset_future
 
     async def _send_loop(self):
         """Send queue handler"""

bellows/zigbee/application.py

@@ -4,8 +4,14 @@
 import logging
 import os
 import statistics
+import sys
 from typing import Dict, Optional
 
+if sys.version_info[:2] < (3, 11):
+    from async_timeout import timeout as asyncio_timeout  # pragma: no cover
+else:
+    from asyncio import timeout as asyncio_timeout  # pragma: no cover
+
 import zigpy.application
 import zigpy.config
 import zigpy.device
@@ -816,9 +822,8 @@ async def send_packet(self, packet: zigpy.types.ZigbeePacket) -> None:
                     return
 
                 # Wait for `messageSentHandler` message
-                send_status, _ = await asyncio.wait_for(
-                    req.result, timeout=APS_ACK_TIMEOUT
-                )
+                async with asyncio_timeout(APS_ACK_TIMEOUT):
+                    send_status, _ = await req.result
 
                 if send_status != t.EmberStatus.SUCCESS:
                     raise zigpy.exceptions.DeliveryError(
@@ -880,9 +885,9 @@ async def _watchdog(self):
         await asyncio.sleep(WATCHDOG_WAKE_PERIOD)
         while True:
             try:
-                await asyncio.wait_for(
-                    self.controller_event.wait(), timeout=WATCHDOG_WAKE_PERIOD * 2
-                )
+                async with asyncio_timeout(WATCHDOG_WAKE_PERIOD * 2):
+                    await self.controller_event.wait()
+
                 if self._ezsp.ezsp_version == 4:
                     await self._ezsp.nop()
                 else:

setup.py

@@ -20,6 +20,7 @@
         "pure_pcapy3==1.0.1",
         "voluptuous",
         "zigpy>=0.54.0",
+        'async-timeout; python_version<"3.11"',
     ],
     dependency_links=[
         "https://codeload.github.com/rcloran/pure-pcapy-3/zip/master",

tests/test_application.py

@@ -60,7 +60,7 @@ def inner(config):
         app = bellows.zigbee.application.ControllerApplication(app_cfg)
 
         app._ezsp = ezsp_mock
-        monkeypatch.setattr(bellows.zigbee.application, "APS_ACK_TIMEOUT", 0.01)
+        monkeypatch.setattr(bellows.zigbee.application, "APS_ACK_TIMEOUT", 0.05)
         app._ctrl_event.set()
         app._in_flight_msg = asyncio.Semaphore()
         app.handle_message = MagicMock()

tests/test_ezsp.py

@@ -70,21 +70,21 @@ def test_attr(ezsp_f):
     assert callable(m)
 
 
-def test_non_existent_attr(ezsp_f):
+async def test_non_existent_attr(ezsp_f):
     with pytest.raises(AttributeError):
-        ezsp_f.nonexistentMethod()
+        await ezsp_f.nonexistentMethod()
 
 
-def test_command(ezsp_f):
+async def test_command(ezsp_f):
     ezsp_f.start_ezsp()
     with patch.object(ezsp_f._protocol, "command") as cmd_mock:
-        ezsp_f.nop()
+        await ezsp_f.nop()
     assert cmd_mock.call_count == 1
 
 
-def test_command_ezsp_stopped(ezsp_f):
+async def test_command_ezsp_stopped(ezsp_f):
     with pytest.raises(EzspError):
-        ezsp_f._command("version")
+        await ezsp_f._command("version")
 
 
 async def _test_list_command(ezsp_f, mockcommand):

tests/test_ezsp_protocol.py

@@ -29,7 +29,10 @@ def prot_hndl():
 
 async def test_command(prot_hndl):
     coro = prot_hndl.command("nop")
-    prot_hndl._awaiting[prot_hndl._seq - 1][2].set_result(True)
+    asyncio.get_running_loop().call_soon(
+        lambda: prot_hndl._awaiting[prot_hndl._seq - 1][2].set_result(True)
+    )
+
     await coro
     assert prot_hndl._gw.data.call_count == 1
 

tests/test_thread.py

@@ -3,6 +3,11 @@
 import threading
 from unittest import mock
 
+if sys.version_info[:2] < (3, 11):
+    from async_timeout import timeout as asyncio_timeout  # pragma: no cover
+else:
+    from asyncio import timeout as asyncio_timeout  # pragma: no cover
+
 import pytest
 
 from bellows.thread import EventLoopThread, ThreadsafeProxy
@@ -48,7 +53,8 @@ async def thread():
     yield thread
     thread.force_stop()
     if thread.thread_complete is not None:
-        await asyncio.wait_for(thread.thread_complete, 1)
+        async with asyncio_timeout(1):
+            await thread.thread_complete
     [t.join(1) for t in threading.enumerate() if "bellows" in t.name]
     threads = [t for t in threading.enumerate() if "bellows" in t.name]
     assert len(threads) == 0
@@ -178,4 +184,5 @@ async def wait_forever():
     # The cancellation should propagate to the outer event loop
     with pytest.raises(asyncio.CancelledError):
         # This will stall forever without the patch
-        await asyncio.wait_for(proxy.wait_forever(), 1)
+        async with asyncio_timeout(1):
+            await proxy.wait_forever()