diff --git a/aioesphomeapi/_frame_helper.py b/aioesphomeapi/_frame_helper.py index 19a2049..327bb08 100644 --- a/aioesphomeapi/_frame_helper.py +++ b/aioesphomeapi/_frame_helper.py @@ -5,6 +5,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from typing import Optional +import async_timeout from noise.connection import NoiseConnection # type: ignore from .core import ( @@ -239,7 +240,8 @@ class APINoiseFrameHelper(APIFrameHelper): async def perform_handshake(self, expected_name: Optional[str]) -> None: # Allow up to 60 seconds for handhsake try: - await asyncio.wait_for(self._perform_handshake(expected_name), timeout=60.0) + async with async_timeout.timeout(60.0): + await self._perform_handshake(expected_name) except asyncio.TimeoutError as err: raise HandshakeAPIError("Timeout during handshake") from err diff --git a/aioesphomeapi/connection.py b/aioesphomeapi/connection.py index 2d1c726..bf565c8 100644 --- a/aioesphomeapi/connection.py +++ b/aioesphomeapi/connection.py @@ -7,6 +7,7 @@ from contextlib import suppress from dataclasses import astuple, dataclass from typing import Any, Callable, Coroutine, List, Optional +import async_timeout from google.protobuf import message import aioesphomeapi.host_resolver as hr @@ -133,7 +134,8 @@ class APIConnection: self._params.port, self._params.zeroconf_instance, ) - return await asyncio.wait_for(coro, 30.0) + async with async_timeout.timeout(30.0): + return await coro except asyncio.TimeoutError as err: raise ResolveAPIError( f"Timeout while resolving IP address for {self.log_name}" @@ -158,7 +160,8 @@ class APIConnection: try: coro = asyncio.get_event_loop().sock_connect(self._socket, sockaddr) - await asyncio.wait_for(coro, 30.0) + async with async_timeout.timeout(30.0): + await coro except OSError as err: raise SocketAPIError(f"Error connecting to {sockaddr}: {err}") from err except asyncio.TimeoutError as err: @@ -229,9 +232,8 @@ class APIConnection: # Wait for keepalive seconds, or ping stop event, whichever happens first try: - await asyncio.wait_for( - self._ping_stop_event.wait(), self._params.keepalive - ) + async with async_timeout.timeout(self._params.keepalive): + await self._ping_stop_event.wait() except asyncio.TimeoutError: pass @@ -279,7 +281,8 @@ class APIConnection: # Allow 2 minutes for connect; this is only as a last measure # to protect from issues if some part of the connect process mistakenly # does not have a timeout - await asyncio.wait_for(_do_connect(), timeout=120.0) + async with async_timeout.timeout(120.0): + await _do_connect() except Exception: # pylint: disable=broad-except # Always clean up the connection if an error occured during connect self._connection_state = ConnectionState.CLOSED @@ -406,7 +409,8 @@ class APIConnection: await self.send_message(send_msg) try: - await asyncio.wait_for(fut, timeout) + async with async_timeout.timeout(timeout): + await fut except asyncio.TimeoutError as err: raise TimeoutAPIError( f"Timeout waiting for response for {type(send_msg)}"