From 7814ce06f40f4b5862206a3ea5fd6bdefb8f9800 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Fri, 21 Jul 2023 13:44:13 -0500
Subject: [PATCH] Fix ESPHome bluetooth client cancel behavior when device
 unexpectedly disconnects (#96918)

---
 .../components/esphome/bluetooth/client.py    | 33 ++++++-------------
 .../components/esphome/manifest.json          |  1 +
 requirements_all.txt                          |  3 ++
 requirements_test_all.txt                     |  3 ++
 4 files changed, 17 insertions(+), 23 deletions(-)

diff --git a/homeassistant/components/esphome/bluetooth/client.py b/homeassistant/components/esphome/bluetooth/client.py
index f7c9da48883..35e66ea7e47 100644
--- a/homeassistant/components/esphome/bluetooth/client.py
+++ b/homeassistant/components/esphome/bluetooth/client.py
@@ -4,7 +4,6 @@ from __future__ import annotations
 import asyncio
 from collections.abc import Callable, Coroutine
 import contextlib
-from functools import partial
 import logging
 from typing import Any, TypeVar, cast
 import uuid
@@ -17,6 +16,7 @@ from aioesphomeapi import (
 )
 from aioesphomeapi.connection import APIConnectionError, TimeoutAPIError
 from aioesphomeapi.core import BluetoothGATTAPIError
+from async_interrupt import interrupt
 import async_timeout
 from bleak.backends.characteristic import BleakGATTCharacteristic
 from bleak.backends.client import BaseBleakClient, NotifyCallback
@@ -57,11 +57,6 @@ def mac_to_int(address: str) -> int:
     return int(address.replace(":", ""), 16)
 
 
-def _on_disconnected(task: asyncio.Task[Any], _: asyncio.Future[None]) -> None:
-    if task and not task.done():
-        task.cancel()
-
-
 def verify_connected(func: _WrapFuncType) -> _WrapFuncType:
     """Define a wrapper throw BleakError if not connected."""
 
@@ -72,25 +67,17 @@ def verify_connected(func: _WrapFuncType) -> _WrapFuncType:
         loop = self._loop
         disconnected_futures = self._disconnected_futures
         disconnected_future = loop.create_future()
-        disconnect_handler = partial(_on_disconnected, asyncio.current_task(loop))
-        disconnected_future.add_done_callback(disconnect_handler)
         disconnected_futures.add(disconnected_future)
+        ble_device = self._ble_device
+        disconnect_message = (
+            f"{self._source_name }: {ble_device.name} - {ble_device.address}: "
+            "Disconnected during operation"
+        )
         try:
-            return await func(self, *args, **kwargs)
-        except asyncio.CancelledError as ex:
-            if not disconnected_future.done():
-                # If the disconnected future is not done, the task was cancelled
-                # externally and we need to raise cancelled error to avoid
-                # blocking the cancellation.
-                raise
-            ble_device = self._ble_device
-            raise BleakError(
-                f"{self._source_name }: {ble_device.name} - {ble_device.address}: "
-                "Disconnected during operation"
-            ) from ex
+            async with interrupt(disconnected_future, BleakError, disconnect_message):
+                return await func(self, *args, **kwargs)
         finally:
             disconnected_futures.discard(disconnected_future)
-            disconnected_future.remove_done_callback(disconnect_handler)
 
     return cast(_WrapFuncType, _async_wrap_bluetooth_connected_operation)
 
@@ -340,7 +327,7 @@ class ESPHomeClient(BaseBleakClient):
                         # exception.
                         await connected_future
                 raise
-            except Exception:
+            except Exception as ex:
                 if connected_future.done():
                     with contextlib.suppress(BleakError):
                         # If the connect call throws an exception,
@@ -350,7 +337,7 @@ class ESPHomeClient(BaseBleakClient):
                         # exception from the connect call as it
                         # will be more descriptive.
                         await connected_future
-                connected_future.cancel()
+                connected_future.cancel(f"Unhandled exception in connect call: {ex}")
                 raise
             await connected_future
 
diff --git a/homeassistant/components/esphome/manifest.json b/homeassistant/components/esphome/manifest.json
index 1caf0e01efc..33c43936544 100644
--- a/homeassistant/components/esphome/manifest.json
+++ b/homeassistant/components/esphome/manifest.json
@@ -15,6 +15,7 @@
   "iot_class": "local_push",
   "loggers": ["aioesphomeapi", "noiseprotocol"],
   "requirements": [
+    "async_interrupt==1.1.1",
     "aioesphomeapi==15.1.14",
     "bluetooth-data-tools==1.6.0",
     "esphome-dashboard-api==1.2.3"
diff --git a/requirements_all.txt b/requirements_all.txt
index 3fb2ff29f9a..8f16fbeb22c 100644
--- a/requirements_all.txt
+++ b/requirements_all.txt
@@ -442,6 +442,9 @@ asterisk-mbox==0.5.0
 # homeassistant.components.yeelight
 async-upnp-client==0.33.2
 
+# homeassistant.components.esphome
+async_interrupt==1.1.1
+
 # homeassistant.components.keyboard_remote
 asyncinotify==4.0.2
 
diff --git a/requirements_test_all.txt b/requirements_test_all.txt
index b99ba4b372c..3ff331e2a00 100644
--- a/requirements_test_all.txt
+++ b/requirements_test_all.txt
@@ -396,6 +396,9 @@ arcam-fmj==1.4.0
 # homeassistant.components.yeelight
 async-upnp-client==0.33.2
 
+# homeassistant.components.esphome
+async_interrupt==1.1.1
+
 # homeassistant.components.sleepiq
 asyncsleepiq==1.3.5