From 6ae79524bd8a50de711a48164da4e919776b95c8 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Tue, 25 Jul 2023 12:30:54 -0500
Subject: [PATCH] Add support for bleak 0.21 (#97212)

---
 .../components/bluetooth/wrappers.py          | 24 +++++++-
 tests/components/bluetooth/test_init.py       | 59 +++++++++++++++++++
 2 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/homeassistant/components/bluetooth/wrappers.py b/homeassistant/components/bluetooth/wrappers.py
index 67e401cd40afff..2ae036080f81cc 100644
--- a/homeassistant/components/bluetooth/wrappers.py
+++ b/homeassistant/components/bluetooth/wrappers.py
@@ -6,13 +6,18 @@
 import contextlib
 from dataclasses import dataclass
 from functools import partial
+import inspect
 import logging
 from typing import TYPE_CHECKING, Any, Final
 
 from bleak import BleakClient, BleakError
 from bleak.backends.client import BaseBleakClient, get_platform_client_backend_type
 from bleak.backends.device import BLEDevice
-from bleak.backends.scanner import AdvertisementDataCallback, BaseBleakScanner
+from bleak.backends.scanner import (
+    AdvertisementData,
+    AdvertisementDataCallback,
+    BaseBleakScanner,
+)
 from bleak_retry_connector import (
     NO_RSSI_VALUE,
     ble_device_description,
@@ -58,6 +63,7 @@ def __init__(
         self._detection_cancel: CALLBACK_TYPE | None = None
         self._mapped_filters: dict[str, set[str]] = {}
         self._advertisement_data_callback: AdvertisementDataCallback | None = None
+        self._background_tasks: set[asyncio.Task] = set()
         remapped_kwargs = {
             "detection_callback": detection_callback,
             "service_uuids": service_uuids or [],
@@ -128,12 +134,24 @@ def _setup_detection_callback(self) -> None:
         """Set up the detection callback."""
         if self._advertisement_data_callback is None:
             return
+        callback = self._advertisement_data_callback
         self._cancel_callback()
         super().register_detection_callback(self._advertisement_data_callback)
         assert models.MANAGER is not None
-        assert self._callback is not None
+
+        if not inspect.iscoroutinefunction(callback):
+            detection_callback = callback
+        else:
+
+            def detection_callback(
+                ble_device: BLEDevice, advertisement_data: AdvertisementData
+            ) -> None:
+                task = asyncio.create_task(callback(ble_device, advertisement_data))
+                self._background_tasks.add(task)
+                task.add_done_callback(self._background_tasks.discard)
+
         self._detection_cancel = models.MANAGER.async_register_bleak_callback(
-            self._callback, self._mapped_filters
+            detection_callback, self._mapped_filters
         )
 
     def __del__(self) -> None:
diff --git a/tests/components/bluetooth/test_init.py b/tests/components/bluetooth/test_init.py
index 24f1039175b8c0..21fade843f54c5 100644
--- a/tests/components/bluetooth/test_init.py
+++ b/tests/components/bluetooth/test_init.py
@@ -2386,6 +2386,65 @@ def _device_detected(
         assert len(detected) == 2
 
 
+async def test_wrapped_instance_with_service_uuids_with_coro_callback(
+    hass: HomeAssistant, mock_bleak_scanner_start: MagicMock, enable_bluetooth: None
+) -> None:
+    """Test consumers can use the wrapped instance with a service_uuids list as if it was normal BleakScanner.
+
+    Verify that coro callbacks are supported.
+    """
+    with patch(
+        "homeassistant.components.bluetooth.async_get_bluetooth", return_value=[]
+    ):
+        await async_setup_with_default_adapter(hass)
+
+    with patch.object(hass.config_entries.flow, "async_init"):
+        hass.bus.async_fire(EVENT_HOMEASSISTANT_STARTED)
+        await hass.async_block_till_done()
+
+        detected = []
+
+        async def _device_detected(
+            device: BLEDevice, advertisement_data: AdvertisementData
+        ) -> None:
+            """Handle a detected device."""
+            detected.append((device, advertisement_data))
+
+        switchbot_device = generate_ble_device("44:44:33:11:23:45", "wohand")
+        switchbot_adv = generate_advertisement_data(
+            local_name="wohand",
+            service_uuids=["cba20d00-224d-11e6-9fb8-0002a5d5c51b"],
+            manufacturer_data={89: b"\xd8.\xad\xcd\r\x85"},
+            service_data={"00000d00-0000-1000-8000-00805f9b34fb": b"H\x10c"},
+        )
+        switchbot_adv_2 = generate_advertisement_data(
+            local_name="wohand",
+            service_uuids=["cba20d00-224d-11e6-9fb8-0002a5d5c51b"],
+            manufacturer_data={89: b"\xd8.\xad\xcd\r\x84"},
+            service_data={"00000d00-0000-1000-8000-00805f9b34fb": b"H\x10c"},
+        )
+        empty_device = generate_ble_device("11:22:33:44:55:66", "empty")
+        empty_adv = generate_advertisement_data(local_name="empty")
+
+        assert _get_manager() is not None
+        scanner = HaBleakScannerWrapper(
+            service_uuids=["cba20d00-224d-11e6-9fb8-0002a5d5c51b"]
+        )
+        scanner.register_detection_callback(_device_detected)
+
+        inject_advertisement(hass, switchbot_device, switchbot_adv)
+        inject_advertisement(hass, switchbot_device, switchbot_adv_2)
+
+        await hass.async_block_till_done()
+
+        assert len(detected) == 2
+
+        # The UUIDs list we created in the wrapped scanner with should be respected
+        # and we should not get another callback
+        inject_advertisement(hass, empty_device, empty_adv)
+        assert len(detected) == 2
+
+
 async def test_wrapped_instance_with_broken_callbacks(
     hass: HomeAssistant, mock_bleak_scanner_start: MagicMock, enable_bluetooth: None
 ) -> None: