Ensure bleak_retry_connector uses HaBleakClientWrapper (#79132)

homeassistant/components/bluetooth/usage.py

@@ -3,20 +3,39 @@
 from __future__ import annotations
 
 import bleak
+from bleak.backends.service import BleakGATTServiceCollection
+import bleak_retry_connector
 
 from .models import HaBleakClientWrapper, HaBleakScannerWrapper
 
 ORIGINAL_BLEAK_SCANNER = bleak.BleakScanner
 ORIGINAL_BLEAK_CLIENT = bleak.BleakClient
+ORIGINAL_BLEAK_RETRY_CONNECTOR_CLIENT = (
+    bleak_retry_connector.BleakClientWithServiceCache
+)
 
 
 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]
+    bleak_retry_connector.BleakClientWithServiceCache = HaBleakClientWithServiceCache  # type: ignore[misc,assignment]
 
 
 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]
+    bleak_retry_connector.BleakClientWithServiceCache = ORIGINAL_BLEAK_RETRY_CONNECTOR_CLIENT  # type: ignore[misc]
+
+
+class HaBleakClientWithServiceCache(HaBleakClientWrapper):
+    """A BleakClient that implements service caching."""
+
+    def set_cached_services(self, services: BleakGATTServiceCollection | None) -> None:
+        """Set the cached services.
+
+        No longer used since bleak 0.17+ has service caching built-in.
+
+        This was only kept for backwards compatibility.
+        """

tests/components/bluetooth/test_usage.py

@@ -5,6 +5,7 @@ from unittest.mock import patch
 
 import bleak
 from bleak.backends.device import BLEDevice
+import bleak_retry_connector
 
 from homeassistant.components.bluetooth.models import (
     HaBleakClientWrapper,
@@ -75,3 +76,32 @@ async def test_bleak_client_reports_with_address(hass, enable_bluetooth, caplog)
 
     assert not isinstance(instance, HaBleakClientWrapper)
     assert "BleakClient with an address instead of a BLEDevice" not in caplog.text
+
+
+async def test_bleak_retry_connector_client_reports_with_address(
+    hass, enable_bluetooth, caplog
+):
+    """Test we report when we pass an address to BleakClientWithServiceCache."""
+    install_multiple_bleak_catcher()
+
+    with patch.object(
+        _get_manager(),
+        "async_ble_device_from_address",
+        return_value=MOCK_BLE_DEVICE,
+    ):
+        instance = bleak_retry_connector.BleakClientWithServiceCache(
+            "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_retry_connector.BleakClientWithServiceCache("00:00:00:00:00:00")
+
+    assert not isinstance(instance, HaBleakClientWrapper)
+    assert "BleakClient with an address instead of a BLEDevice" not in caplog.text