Prevent bluetooth scanner from being shutdown by BleakClient not using BLEDevice (#76945)

homeassistant/components/bluetooth/models.py

@@ -9,6 +9,8 @@ from enum import Enum
 import logging
 from typing import TYPE_CHECKING, Any, Final
 
+from bleak import BleakClient, BleakError
+from bleak.backends.device import BLEDevice
 from bleak.backends.scanner import (
     AdvertisementData,
     AdvertisementDataCallback,
@@ -16,10 +18,10 @@ from bleak.backends.scanner import (
 )
 
 from homeassistant.core import CALLBACK_TYPE
+from homeassistant.helpers.frame import report
 from homeassistant.helpers.service_info.bluetooth import BluetoothServiceInfo
 
 if TYPE_CHECKING:
-    from bleak.backends.device import BLEDevice
 
     from .manager import BluetoothManager
 
@@ -165,3 +167,33 @@ class HaBleakScannerWrapper(BaseBleakScanner):
             # Nothing to do if event loop is already closed
             with contextlib.suppress(RuntimeError):
                 asyncio.get_running_loop().call_soon_threadsafe(self._detection_cancel)
+
+
+class HaBleakClientWrapper(BleakClient):
+    """Wrap the BleakClient to ensure it does not shutdown our scanner.
+
+    If an address is passed into BleakClient instead of a BLEDevice,
+    bleak will quietly start a new scanner under the hood to resolve
+    the address. This can cause a conflict with our scanner. We need
+    to handle translating the address to the BLEDevice in this case
+    to avoid the whole stack from getting stuck in an in progress state
+    when an integration does this.
+    """
+
+    def __init__(
+        self, address_or_ble_device: str | BLEDevice, *args: Any, **kwargs: Any
+    ) -> None:
+        """Initialize the BleakClient."""
+        if isinstance(address_or_ble_device, BLEDevice):
+            super().__init__(address_or_ble_device, *args, **kwargs)
+            return
+        report(
+            "attempted to call BleakClient with an address instead of a BLEDevice",
+            exclude_integrations={"bluetooth"},
+            error_if_core=False,
+        )
+        assert MANAGER is not None
+        ble_device = MANAGER.async_ble_device_from_address(address_or_ble_device)
+        if ble_device is None:
+            raise BleakError(f"No device found for address {address_or_ble_device}")
+        super().__init__(ble_device, *args, **kwargs)

homeassistant/components/bluetooth/usage.py

@@ -4,16 +4,19 @@ from __future__ import annotations
 
 import bleak
 
-from .models import HaBleakScannerWrapper
+from .models import HaBleakClientWrapper, HaBleakScannerWrapper
 
 ORIGINAL_BLEAK_SCANNER = bleak.BleakScanner
+ORIGINAL_BLEAK_CLIENT = bleak.BleakClient
 
 
 def install_multiple_bleak_catcher() -> None:
     """Wrap the bleak classes to return the shared instance if multiple instances are detected."""
     bleak.BleakScanner = HaBleakScannerWrapper  # type: ignore[misc, assignment]
+    bleak.BleakClient = HaBleakClientWrapper  # type: ignore[misc]
 
 
 def uninstall_multiple_bleak_catcher() -> None:
     """Unwrap the bleak classes."""
     bleak.BleakScanner = ORIGINAL_BLEAK_SCANNER  # type: ignore[misc]
+    bleak.BleakClient = ORIGINAL_BLEAK_CLIENT  # type: ignore[misc]

tests/components/bluetooth/test_usage.py

@@ -1,14 +1,27 @@
 """Tests for the Bluetooth integration."""
 
 
-import bleak
+from unittest.mock import patch
 
-from homeassistant.components.bluetooth.models import HaBleakScannerWrapper
+import bleak
+from bleak.backends.device import BLEDevice
+import pytest
+
+from homeassistant.components.bluetooth.models import (
+    HaBleakClientWrapper,
+    HaBleakScannerWrapper,
+)
 from homeassistant.components.bluetooth.usage import (
     install_multiple_bleak_catcher,
     uninstall_multiple_bleak_catcher,
 )
 
+from . import _get_manager
+
+MOCK_BLE_DEVICE = BLEDevice(
+    "00:00:00:00:00:00", "any", delegate="", details={"path": "/dev/hci0/device"}
+)
+
 
 async def test_multiple_bleak_scanner_instances(hass):
     """Test creating multiple BleakScanners without an integration."""
@@ -23,3 +36,46 @@ async def test_multiple_bleak_scanner_instances(hass):
     instance = bleak.BleakScanner()
 
     assert not isinstance(instance, HaBleakScannerWrapper)
+
+
+async def test_wrapping_bleak_client(hass, enable_bluetooth):
+    """Test we wrap BleakClient."""
+    install_multiple_bleak_catcher()
+
+    instance = bleak.BleakClient(MOCK_BLE_DEVICE)
+
+    assert isinstance(instance, HaBleakClientWrapper)
+
+    uninstall_multiple_bleak_catcher()
+
+    instance = bleak.BleakClient(MOCK_BLE_DEVICE)
+
+    assert not isinstance(instance, HaBleakClientWrapper)
+
+
+async def test_bleak_client_reports_with_address(hass, enable_bluetooth, caplog):
+    """Test we report when we pass an address to BleakClient."""
+    install_multiple_bleak_catcher()
+
+    with pytest.raises(bleak.BleakError):
+        instance = bleak.BleakClient("00:00:00:00:00:00")
+
+    with patch.object(
+        _get_manager(),
+        "async_ble_device_from_address",
+        return_value=MOCK_BLE_DEVICE,
+    ):
+        instance = bleak.BleakClient("00:00:00:00:00:00")
+
+    assert "BleakClient with an address instead of a BLEDevice" in caplog.text
+
+    assert isinstance(instance, HaBleakClientWrapper)
+
+    uninstall_multiple_bleak_catcher()
+
+    caplog.clear()
+
+    instance = bleak.BleakClient("00:00:00:00:00:00")
+
+    assert not isinstance(instance, HaBleakClientWrapper)
+    assert "BleakClient with an address instead of a BLEDevice" not in caplog.text