diff --git a/homeassistant/components/esphome/bluetooth/client.py b/homeassistant/components/esphome/bluetooth/client.py index d452ab8764a..00b9883f261 100644 --- a/homeassistant/components/esphome/bluetooth/client.py +++ b/homeassistant/components/esphome/bluetooth/client.py @@ -62,29 +62,32 @@ def verify_connected(func: _WrapFuncType) -> _WrapFuncType: async def _async_wrap_bluetooth_connected_operation( self: ESPHomeClient, *args: Any, **kwargs: Any ) -> Any: - disconnected_event = ( - self._disconnected_event # pylint: disable=protected-access + loop = self._loop # pylint: disable=protected-access + disconnected_futures = ( + self._disconnected_futures # pylint: disable=protected-access ) - if not disconnected_event: - raise BleakError("Not connected") - action_task = asyncio.create_task(func(self, *args, **kwargs)) - disconnect_task = asyncio.create_task(disconnected_event.wait()) - await asyncio.wait( - (action_task, disconnect_task), - return_when=asyncio.FIRST_COMPLETED, - ) - if disconnect_task.done(): - action_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await action_task + disconnected_future = loop.create_future() + disconnected_futures.add(disconnected_future) + task = asyncio.current_task(loop) + + def _on_disconnected(fut: asyncio.Future[None]) -> None: + if task and not task.done(): + task.cancel() + + disconnected_future.add_done_callback(_on_disconnected) + try: + return await func(self, *args, **kwargs) + except asyncio.CancelledError as ex: + source_name = self._source_name # pylint: disable=protected-access + ble_device = self._ble_device # pylint: disable=protected-access raise BleakError( - f"{self._source_name}: " # pylint: disable=protected-access - f"{self._ble_device.name} - " # pylint: disable=protected-access - f" {self._ble_device.address}: " # pylint: disable=protected-access + f"{source_name}: {ble_device.name} - {ble_device.address}: " "Disconnected during operation" - ) - return action_task.result() + ) from ex + finally: + disconnected_futures.discard(disconnected_future) + disconnected_future.remove_done_callback(_on_disconnected) return cast(_WrapFuncType, _async_wrap_bluetooth_connected_operation) @@ -152,7 +155,8 @@ class ESPHomeClient(BaseBleakClient): self._notify_cancels: dict[ int, tuple[Callable[[], Coroutine[Any, Any, None]], Callable[[], None]] ] = {} - self._disconnected_event: asyncio.Event | None = None + self._loop = asyncio.get_running_loop() + self._disconnected_futures: set[asyncio.Future[None]] = set() device_info = self.entry_data.device_info assert device_info is not None self._device_info = device_info @@ -192,9 +196,10 @@ class ESPHomeClient(BaseBleakClient): for _, notify_abort in self._notify_cancels.values(): notify_abort() self._notify_cancels.clear() - if self._disconnected_event: - self._disconnected_event.set() - self._disconnected_event = None + for future in self._disconnected_futures: + if not future.done(): + future.set_result(None) + self._disconnected_futures.clear() self._unsubscribe_connection_state() def _async_ble_device_disconnected(self) -> None: @@ -359,7 +364,6 @@ class ESPHomeClient(BaseBleakClient): await self.disconnect() raise - self._disconnected_event = asyncio.Event() return True @api_error_as_bleak_error