feat: add support for getting callbacks when adapter allocations change (#115)

Author: bdraco

src/habluetooth/manager.pxd

@@ -51,6 +51,8 @@ cdef class BluetoothManager:
     cdef public object _adapter_refresh_future
     cdef public object _recovery_lock
     cdef public set _disappeared_callbacks
+    cdef public set _allocations_callbacks
+    cdef public object _cancel_allocation_callbacks
 
     @cython.locals(stale_seconds=float)
     cdef bint _prefer_previous_adv_from_different_source(

src/habluetooth/manager.py

@@ -10,7 +10,12 @@
 from typing import TYPE_CHECKING, Any, Final
 
 from bleak.backends.scanner import AdvertisementDataCallback
-from bleak_retry_connector import NO_RSSI_VALUE, BleakSlotManager
+from bleak_retry_connector import (
+    NO_RSSI_VALUE,
+    AllocationChangeEvent,
+    Allocations,
+    BleakSlotManager,
+)
 from bluetooth_adapters import (
     ADAPTER_ADDRESS,
     ADAPTER_PASSIVE_SCAN,
@@ -99,8 +104,10 @@ class BluetoothManager:
         "_adapters",
         "_advertisement_tracker",
         "_all_history",
+        "_allocations_callbacks",
         "_bleak_callbacks",
         "_bluetooth_adapters",
+        "_cancel_allocation_callbacks",
         "_cancel_unavailable_tracking",
         "_connectable_history",
         "_connectable_scanners",
@@ -146,12 +153,18 @@ def __init__(
         self._sources: dict[str, BaseHaScanner] = {}
         self._bluetooth_adapters = bluetooth_adapters
         self.slot_manager = slot_manager
+        self._cancel_allocation_callbacks = (
+            self.slot_manager.register_allocation_callback(
+                self._async_slot_manager_changed
+            )
+        )
         self._debug = _LOGGER.isEnabledFor(logging.DEBUG)
         self.shutdown = False
         self._loop: asyncio.AbstractEventLoop | None = None
         self._adapter_refresh_future: asyncio.Future[None] | None = None
         self._recovery_lock: asyncio.Lock = asyncio.Lock()
         self._disappeared_callbacks: set[Callable[[str], None]] = set()
+        self._allocations_callbacks: set[Callable[[Allocations], None]] = set()
 
     @property
     def supports_passive_scan(self) -> bool:
@@ -203,13 +216,7 @@ def async_register_disappeared_callback(
     ) -> CALLBACK_TYPE:
         """Register a callback to be called when an address disappears."""
         self._disappeared_callbacks.add(callback)
-        return partial(self._async_remove_disappeared_callback, callback)
-
-    def _async_remove_disappeared_callback(
-        self, callback: Callable[[str], None]
-    ) -> None:
-        """Remove a disappeared callback."""
-        self._disappeared_callbacks.discard(callback)
+        return partial(self._disappeared_callbacks.discard, callback)
 
     async def _async_refresh_adapters(self) -> None:
         """Refresh the adapters."""
@@ -283,6 +290,7 @@ def async_stop(self) -> None:
             self._cancel_unavailable_tracking.cancel()
             self._cancel_unavailable_tracking = None
         uninstall_multiple_bleak_catcher()
+        self._cancel_allocation_callbacks()
 
     def async_scanner_devices_by_address(
         self, address: str, connectable: bool
@@ -749,3 +757,24 @@ def async_set_fallback_availability_interval(
     ) -> None:
         """Override the fallback availability timeout for a MAC address."""
         self._fallback_intervals[address] = interval
+
+    def _async_slot_manager_changed(self, event: AllocationChangeEvent) -> None:
+        """Handle slot manager changes."""
+        self.async_on_allocation_changed(
+            self.slot_manager.get_allocations(event.adapter)
+        )
+
+    def async_on_allocation_changed(self, allocations: Allocations) -> None:
+        """Call allocation callbacks."""
+        for callback_ in self._allocations_callbacks:
+            try:
+                callback_(allocations)
+            except Exception:
+                _LOGGER.exception("Error in allocation callback")
+
+    def async_register_allocation_callback(
+        self, callback: Callable[[str], None]
+    ) -> CALLBACK_TYPE:
+        """Register a callback to be called when an allocations change."""
+        self._allocations_callbacks.add(callback)
+        return partial(self._allocations_callbacks.discard, callback)

tests/conftest.py

@@ -57,7 +57,9 @@ def discovered_devices_and_advertisement_data(
 def manager():
     slot_manager = BleakSlotManager()
     bluetooth_adapters = FakeBluetoothAdapters()
-    set_manager(BluetoothManager(bluetooth_adapters, slot_manager))
+    manager = BluetoothManager(bluetooth_adapters, slot_manager)
+    set_manager(manager)
+    manager.async_stop()
 
 
 @pytest_asyncio.fixture(name="enable_bluetooth")

tests/test_manager.py

@@ -6,7 +6,7 @@
 from unittest.mock import patch
 
 import pytest
-from bleak_retry_connector import BleakSlotManager
+from bleak_retry_connector import AllocationChange, Allocations, BleakSlotManager
 from bluetooth_adapters.systems.linux import LinuxAdapters
 from freezegun import freeze_time
 
@@ -208,3 +208,87 @@ def _ok_callback(_address: str) -> None:
 
     cancel1()
     cancel2()
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("enable_bluetooth")
+async def test_async_register_allocation_callback(
+    register_hci0_scanner: None,
+    register_hci1_scanner: None,
+) -> None:
+    """Test bluetooth async_register_allocation_callback handles failures."""
+    manager = get_manager()
+    assert manager._loop is not None
+
+    address = "44:44:33:11:23:12"
+
+    switchbot_device_signal_100 = generate_ble_device(
+        address, "wohand_signal_100", rssi=-100
+    )
+    switchbot_adv_signal_100 = generate_advertisement_data(
+        local_name="wohand_signal_100", service_uuids=[]
+    )
+    inject_advertisement_with_source(
+        switchbot_device_signal_100, switchbot_adv_signal_100, "hci0"
+    )
+
+    failed_allocations: list[Allocations] = []
+
+    def _failing_callback(allocations: Allocations) -> None:
+        """Failing callback."""
+        failed_allocations.append(allocations)
+        raise ValueError("This is a test")
+
+    ok_allocations: list[Allocations] = []
+
+    def _ok_callback(allocations: Allocations) -> None:
+        """Ok callback."""
+        ok_allocations.append(allocations)
+
+    cancel1 = manager.async_register_allocation_callback(_failing_callback)
+    # Make sure the second callback still works if the first one fails and
+    # raises an exception
+    cancel2 = manager.async_register_allocation_callback(_ok_callback)
+
+    switchbot_adv_signal_100 = generate_advertisement_data(
+        local_name="wohand_signal_100",
+        manufacturer_data={123: b"abc"},
+        service_uuids=[],
+        rssi=-80,
+    )
+    inject_advertisement_with_source(
+        switchbot_device_signal_100, switchbot_adv_signal_100, "hci1"
+    )
+
+    manager.async_on_allocation_changed(
+        Allocations(
+            "hci0",
+            5,
+            4,
+            ["44:44:33:11:23:12"],
+        )
+    )
+
+    assert len(ok_allocations) == 1
+    assert ok_allocations[0] == Allocations(
+        "hci0",
+        5,
+        4,
+        ["44:44:33:11:23:12"],
+    )
+    assert len(failed_allocations) == 1
+    assert failed_allocations[0] == Allocations(
+        "hci0",
+        5,
+        4,
+        ["44:44:33:11:23:12"],
+    )
+
+    manager.slot_manager._allocations_by_adapter["hci0"] = {}
+    manager.slot_manager._call_callbacks(
+        AllocationChange.ALLOCATED, "/org/bluez/hci0/dev_44_44_33_11_23_12"
+    )
+    assert len(ok_allocations) == 2
+
+    cancel1()
+    cancel2()