From 7b8ed9fd0388dc8510f5b95681fc67d4f9a093eb Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 7 May 2023 09:05:51 -0500
Subject: [PATCH] Fix BLE connection timing out on failure instead of raising
 (#429)

---
 aioesphomeapi/client.py | 54 +++++++++++++++++++++++++----------------
 1 file changed, 33 insertions(+), 21 deletions(-)

diff --git a/aioesphomeapi/client.py b/aioesphomeapi/client.py
index 709e790..940233c 100644
--- a/aioesphomeapi/client.py
+++ b/aioesphomeapi/client.py
@@ -529,8 +529,10 @@ class APIClient:
             resp = BluetoothDeviceConnection.from_pb(msg)
             if address == resp.address:
                 on_bluetooth_connection_state(resp.connected, resp.mtu, resp.error)
-                if resp.connected:
-                    event.set()
+                # Resolve on ANY connection state since we do not want
+                # to wait the whole timeout if the device disconnects
+                # or we get an error.
+                event.set()
 
         assert self._connection is not None
         if has_cache:
@@ -570,26 +572,27 @@ class APIClient:
                 # Disconnect before raising the exception to ensure
                 # the slot is recovered before the timeout is raised
                 # to avoid race were we run out even though we have a slot.
-                await self.bluetooth_device_disconnect(address)
                 addr = to_human_readable_address(address)
                 _LOGGER.debug("%s: Connecting timed out, waiting for disconnect", addr)
+                disconnect_timed_out = False
                 try:
-                    async with async_timeout.timeout(disconnect_timeout):
-                        await event.wait()
-                        disconnect_timed_out = False
-                except asyncio.TimeoutError:
-                    disconnect_timed_out = True
-                _LOGGER.debug(
-                    "%s: Disconnect timed out: %s", addr, disconnect_timed_out
-                )
-                try:
-                    unsub()
-                except (KeyError, ValueError):
-                    _LOGGER.warning(
-                        "%s: Bluetooth device connection timed out but already unsubscribed "
-                        "(likely due to unexpected disconnect)",
-                        addr,
+                    await self.bluetooth_device_disconnect(
+                        address, timeout=disconnect_timeout
                     )
+                except TimeoutAPIError:
+                    disconnect_timed_out = True
+                    _LOGGER.debug(
+                        "%s: Disconnect timed out: %s", addr, disconnect_timed_out
+                    )
+                finally:
+                    try:
+                        unsub()
+                    except (KeyError, ValueError):
+                        _LOGGER.warning(
+                            "%s: Bluetooth device connection timed out but already unsubscribed "
+                            "(likely due to unexpected disconnect)",
+                            addr,
+                        )
                 raise TimeoutAPIError(
                     f"Timeout waiting for connect response while connecting to {addr} "
                     f"after {timeout}s, disconnect timed out: {disconnect_timed_out}, "
@@ -693,15 +696,24 @@ class APIClient:
 
         return BluetoothDeviceClearCache.from_pb(response)
 
-    async def bluetooth_device_disconnect(self, address: int) -> None:
+    async def bluetooth_device_disconnect(
+        self, address: int, timeout: float = DEFAULT_BLE_DISCONNECT_TIMEOUT
+    ) -> None:
         self._check_authenticated()
 
+        def predicate_func(msg: BluetoothDeviceConnectionResponse) -> bool:
+            return bool(msg.address == address and not msg.connected)
+
         assert self._connection is not None
-        self._connection.send_message(
+        await self._connection.send_message_await_response_complex(
             BluetoothDeviceRequest(
                 address=address,
                 request_type=BluetoothDeviceRequestType.DISCONNECT,
-            )
+            ),
+            predicate_func,
+            predicate_func,
+            (BluetoothDeviceConnectionResponse,),
+            timeout=timeout,
         )
 
     async def bluetooth_gatt_get_services(