Replace asyncio.wait_for with async_timeout in baf (#78445)

home-assistant · Sep 14, 2022 · b7e9fcb · b7e9fcb
1 parent db44be7
commit b7e9fcb
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/homeassistant/components/baf/__init__.py b/homeassistant/components/baf/__init__.py
@@ -5,6 +5,7 @@
 
 from aiobafi6 import Device, Service
 from aiobafi6.discovery import PORT
+import async_timeout
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_IP_ADDRESS, Platform
@@ -34,7 +35,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     run_future = device.async_run()
 
     try:
-        await asyncio.wait_for(device.async_wait_available(), timeout=RUN_TIMEOUT)
+        async with async_timeout.timeout(RUN_TIMEOUT):
+            await device.async_wait_available()
     except asyncio.TimeoutError as ex:
         run_future.cancel()
         raise ConfigEntryNotReady(f"Timed out connecting to {ip_address}") from ex

diff --git a/homeassistant/components/baf/config_flow.py b/homeassistant/components/baf/config_flow.py
@@ -7,6 +7,7 @@
 
 from aiobafi6 import Device, Service
 from aiobafi6.discovery import PORT
+import async_timeout
 import voluptuous as vol
 
 from homeassistant import config_entries
@@ -26,7 +27,8 @@ async def async_try_connect(ip_address: str) -> Device:
     device = Device(Service(ip_addresses=[ip_address], port=PORT))
     run_future = device.async_run()
     try:
-        await asyncio.wait_for(device.async_wait_available(), timeout=RUN_TIMEOUT)
+        async with async_timeout.timeout(RUN_TIMEOUT):
+            await device.async_wait_available()
     except asyncio.TimeoutError as ex:
         raise CannotConnect from ex
     finally: