Upstream/aioesphomeapi
1
0
Fork 0
mirror of https://github.com/esphome/aioesphomeapi.git synced 2025-02-17 01:51:23 +01:00
Code Issues Projects Releases Wiki Activity

Change out asyncio.wait_for with async_timeout (#243)

Browse Source
This commit is contained in:
J. Nick Koston 2022-08-21 17:26:53 -10:00 committed by GitHub
parent 88749ad2cd
commit 1ffe252f6f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 14 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
aioesphomeapi/_frame_helper.py
View File

@ -5,6 +5,7 @@ from abc import ABC, abstractmethod
from dataclasses import dataclass from dataclasses import dataclass
from typing import Optional from typing import Optional
import async_timeout
from noise.connection import NoiseConnection # type: ignore from noise.connection import NoiseConnection # type: ignore
from .core import ( from .core import (
@ -239,7 +240,8 @@ class APINoiseFrameHelper(APIFrameHelper):
async def perform_handshake(self, expected_name: Optional[str]) -> None: async def perform_handshake(self, expected_name: Optional[str]) -> None:
# Allow up to 60 seconds for handhsake # Allow up to 60 seconds for handhsake
try: 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: except asyncio.TimeoutError as err:
raise HandshakeAPIError("Timeout during handshake") from err raise HandshakeAPIError("Timeout during handshake") from err

18
aioesphomeapi/connection.py
View File

@ -7,6 +7,7 @@ from contextlib import suppress
from dataclasses import astuple, dataclass from dataclasses import astuple, dataclass
from typing import Any, Callable, Coroutine, List, Optional from typing import Any, Callable, Coroutine, List, Optional
import async_timeout
from google.protobuf import message from google.protobuf import message
import aioesphomeapi.host_resolver as hr import aioesphomeapi.host_resolver as hr
@ -133,7 +134,8 @@ class APIConnection:
self._params.port, self._params.port,
self._params.zeroconf_instance, 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: except asyncio.TimeoutError as err:
raise ResolveAPIError( raise ResolveAPIError(
f"Timeout while resolving IP address for {self.log_name}" f"Timeout while resolving IP address for {self.log_name}"
@ -158,7 +160,8 @@ class APIConnection:
try: try:
coro = asyncio.get_event_loop().sock_connect(self._socket, sockaddr) 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: except OSError as err:
raise SocketAPIError(f"Error connecting to {sockaddr}: {err}") from err raise SocketAPIError(f"Error connecting to {sockaddr}: {err}") from err
except asyncio.TimeoutError as err: except asyncio.TimeoutError as err:
@ -229,9 +232,8 @@ class APIConnection:
# Wait for keepalive seconds, or ping stop event, whichever happens first # Wait for keepalive seconds, or ping stop event, whichever happens first
try: try:
await asyncio.wait_for( async with async_timeout.timeout(self._params.keepalive):
self._ping_stop_event.wait(), self._params.keepalive await self._ping_stop_event.wait()
)
except asyncio.TimeoutError: except asyncio.TimeoutError:
pass pass
@ -279,7 +281,8 @@ class APIConnection:
# Allow 2 minutes for connect; this is only as a last measure # Allow 2 minutes for connect; this is only as a last measure
# to protect from issues if some part of the connect process mistakenly # to protect from issues if some part of the connect process mistakenly
# does not have a timeout # 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 except Exception: # pylint: disable=broad-except
# Always clean up the connection if an error occured during connect # Always clean up the connection if an error occured during connect
self._connection_state = ConnectionState.CLOSED self._connection_state = ConnectionState.CLOSED
@ -406,7 +409,8 @@ class APIConnection:
await self.send_message(send_msg) await self.send_message(send_msg)
try: try:
await asyncio.wait_for(fut, timeout) async with async_timeout.timeout(timeout):
await fut
except asyncio.TimeoutError as err: except asyncio.TimeoutError as err:
raise TimeoutAPIError( raise TimeoutAPIError(
f"Timeout waiting for response for {type(send_msg)}" f"Timeout waiting for response for {type(send_msg)}"