feat: Add API for getting current slot allocations (#116)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

src/habluetooth/__init__.py

@@ -18,6 +18,7 @@
     BluetoothServiceInfo,
     BluetoothServiceInfoBleak,
     HaBluetoothConnector,
+    HaBluetoothSlotAllocations,
 )
 from .scanner import BluetoothScanningMode, HaScanner, ScannerStartError
 from .scanner_device import BluetoothScannerDevice
@@ -41,6 +42,7 @@
     "HaBleakClientWrapper",
     "HaBleakScannerWrapper",
     "HaBluetoothConnector",
+    "HaBluetoothSlotAllocations",
     "HaScanner",
     "ScannerStartError",
     "get_manager",

src/habluetooth/manager.pxd

@@ -51,8 +51,10 @@ 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 dict _allocations_callbacks
     cdef public object _cancel_allocation_callbacks
+    cdef public dict _adapter_sources
+    cdef public dict _allocations
 
     @cython.locals(stale_seconds=float)
     cdef bint _prefer_previous_adv_from_different_source(

src/habluetooth/manager.py

@@ -35,7 +35,7 @@
     FALLBACK_MAXIMUM_STALE_ADVERTISEMENT_SECONDS,
     UNAVAILABLE_TRACK_SECONDS,
 )
-from .models import BluetoothServiceInfoBleak
+from .models import BluetoothServiceInfoBleak, HaBluetoothSlotAllocations
 from .scanner_device import BluetoothScannerDevice
 from .usage import install_multiple_bleak_catcher, uninstall_multiple_bleak_catcher
 from .util import async_reset_adapter
@@ -101,9 +101,11 @@ class BluetoothManager:
 
     __slots__ = (
         "_adapter_refresh_future",
+        "_adapter_sources",
         "_adapters",
         "_advertisement_tracker",
         "_all_history",
+        "_allocations",
         "_allocations_callbacks",
         "_bleak_callbacks",
         "_bluetooth_adapters",
@@ -150,6 +152,8 @@ def __init__(
         self._non_connectable_scanners: set[BaseHaScanner] = set()
         self._connectable_scanners: set[BaseHaScanner] = set()
         self._adapters: dict[str, AdapterDetails] = {}
+        self._adapter_sources: dict[str, str] = {}
+        self._allocations: dict[str, HaBluetoothSlotAllocations] = {}
         self._sources: dict[str, BaseHaScanner] = {}
         self._bluetooth_adapters = bluetooth_adapters
         self.slot_manager = slot_manager
@@ -164,7 +168,9 @@ def __init__(
         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()
+        self._allocations_callbacks: dict[
+            str | None, set[Callable[[HaBluetoothSlotAllocations], None]]
+        ] = {}
 
     @property
     def supports_passive_scan(self) -> bool:
@@ -698,6 +704,8 @@ def _async_unregister_scanner_internal(
         self._advertisement_tracker.async_remove_source(scanner.source)
         scanners.remove(scanner)
         del self._sources[scanner.source]
+        del self._adapter_sources[scanner.adapter]
+        self._allocations.pop(scanner.source, None)
         if connection_slots:
             self.slot_manager.remove_adapter(scanner.adapter)
 
@@ -714,6 +722,7 @@ def async_register_scanner(
             scanners = self._non_connectable_scanners
         scanners.add(scanner)
         self._sources[scanner.source] = scanner
+        self._adapter_sources[scanner.adapter] = scanner.source
         if connection_slots:
             self.slot_manager.register_adapter(scanner.adapter, connection_slots)
         return partial(
@@ -766,15 +775,48 @@ def _async_slot_manager_changed(self, event: AllocationChangeEvent) -> None:
 
     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")
+        source = self._adapter_sources.get(allocations.adapter, allocations.adapter)
+        ha_slot_allocations = HaBluetoothSlotAllocations(
+            source=source,
+            slots=allocations.slots,
+            free=allocations.free,
+            allocated=allocations.allocated,
+        )
+        self._allocations[source] = ha_slot_allocations
+        for source_key in (source, None):
+            if not (
+                allocation_callbacks := self._allocations_callbacks.get(source_key)
+            ):
+                continue
+            for callback_ in allocation_callbacks:
+                try:
+                    callback_(ha_slot_allocations)
+                except Exception:
+                    _LOGGER.exception("Error in allocation callback")
+
+    def async_current_allocations(
+        self, source: str | None = None
+    ) -> list[HaBluetoothSlotAllocations] | None:
+        """Return the current allocations."""
+        if source:
+            if allocations := self._allocations.get(source):
+                return [allocations]
+            return []
+        return list(self._allocations.values())
 
     def async_register_allocation_callback(
-        self, callback: Callable[[str], None]
+        self,
+        callback: Callable[[HaBluetoothSlotAllocations], None],
+        source: str | None = 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)
+        self._allocations_callbacks.setdefault(source, set()).add(callback)
+        return partial(self._async_unregister_allocation_callback, callback, source)
+
+    def _async_unregister_allocation_callback(
+        self, callback: Callable[[HaBluetoothSlotAllocations], None], source: str | None
+    ) -> None:
+        if (callbacks := self._allocations_callbacks.get(source)) is not None:
+            callbacks.discard(callback)
+            if not callbacks:
+                del self._allocations_callbacks[source]

src/habluetooth/models.py

@@ -47,6 +47,16 @@ def set_manager(manager: BluetoothManager) -> None:
     CentralBluetoothManager.manager = manager
 
 
+@dataclass(slots=True)
+class HaBluetoothSlotAllocations:
+    """Data for how to allocate slots for BLEDevice connections."""
+
+    source: str  # Adapter MAC
+    slots: int  # Number of slots
+    free: int  # Number of free slots
+    allocated: list[str]  # Addresses of connected devices
+
+
 @dataclass(slots=True)
 class HaBluetoothConnector:
     """Data for how to connect a BLEDevice from a given scanner."""

tests/conftest.py

@@ -53,12 +53,14 @@ def discovered_devices_and_advertisement_data(
         return {}
 
 
-@pytest.fixture(scope="session", autouse=True)
-def manager():
+@pytest_asyncio.fixture(scope="session", autouse=True)
+async def manager() -> AsyncGenerator[None, None]:
     slot_manager = BleakSlotManager()
     bluetooth_adapters = FakeBluetoothAdapters()
     manager = BluetoothManager(bluetooth_adapters, slot_manager)
     set_manager(manager)
+    await manager.async_setup()
+    yield
     manager.async_stop()
 
 
@@ -158,18 +160,18 @@ def two_adapters_fixture():
 @pytest.fixture
 def register_hci0_scanner() -> Generator[None, None, None]:
     """Register an hci0 scanner."""
-    hci0_scanner = FakeScanner("hci0", "hci0")
+    hci0_scanner = FakeScanner("AA:BB:CC:DD:EE:00", "hci0")
     manager = get_manager()
-    cancel = manager.async_register_scanner(hci0_scanner)
+    cancel = manager.async_register_scanner(hci0_scanner, connection_slots=5)
     yield
     cancel()
 
 
 @pytest.fixture
 def register_hci1_scanner() -> Generator[None, None, None]:
     """Register an hci1 scanner."""
-    hci1_scanner = FakeScanner("hci1", "hci1")
+    hci1_scanner = FakeScanner("AA:BB:CC:DD:EE:11", "hci1")
     manager = get_manager()
-    cancel = manager.async_register_scanner(hci1_scanner)
+    cancel = manager.async_register_scanner(hci1_scanner, connection_slots=5)
     yield
     cancel()

tests/test_manager.py

@@ -10,7 +10,12 @@
 from bluetooth_adapters.systems.linux import LinuxAdapters
 from freezegun import freeze_time
 
-from habluetooth import BluetoothManager, get_manager, set_manager
+from habluetooth import (
+    BluetoothManager,
+    HaBluetoothSlotAllocations,
+    get_manager,
+    set_manager,
+)
 
 from . import (
     async_fire_time_changed,
@@ -232,16 +237,16 @@ async def test_async_register_allocation_callback(
         switchbot_device_signal_100, switchbot_adv_signal_100, "hci0"
     )
 
-    failed_allocations: list[Allocations] = []
+    failed_allocations: list[HaBluetoothSlotAllocations] = []
 
-    def _failing_callback(allocations: Allocations) -> None:
+    def _failing_callback(allocations: HaBluetoothSlotAllocations) -> None:
         """Failing callback."""
         failed_allocations.append(allocations)
         raise ValueError("This is a test")
 
-    ok_allocations: list[Allocations] = []
+    ok_allocations: list[HaBluetoothSlotAllocations] = []
 
-    def _ok_callback(allocations: Allocations) -> None:
+    def _ok_callback(allocations: HaBluetoothSlotAllocations) -> None:
         """Ok callback."""
         ok_allocations.append(allocations)
 
@@ -260,35 +265,52 @@ def _ok_callback(allocations: Allocations) -> None:
         switchbot_device_signal_100, switchbot_adv_signal_100, "hci1"
     )
 
+    assert manager.async_current_allocations() == []
     manager.async_on_allocation_changed(
         Allocations(
-            "hci0",
+            "AA:BB:CC:DD:EE:00",
             5,
             4,
             ["44:44:33:11:23:12"],
         )
     )
 
     assert len(ok_allocations) == 1
-    assert ok_allocations[0] == Allocations(
-        "hci0",
+    assert ok_allocations[0] == HaBluetoothSlotAllocations(
+        "AA:BB:CC:DD:EE:00",
         5,
         4,
         ["44:44:33:11:23:12"],
     )
     assert len(failed_allocations) == 1
-    assert failed_allocations[0] == Allocations(
-        "hci0",
+    assert failed_allocations[0] == HaBluetoothSlotAllocations(
+        "AA:BB:CC:DD:EE:00",
         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"
-    )
+    with patch.object(
+        manager.slot_manager,
+        "get_allocations",
+        return_value=Allocations(
+            adapter="hci0",
+            slots=5,
+            free=4,
+            allocated=["44:44:33:11:23:12"],
+        ),
+    ):
+        manager.slot_manager._call_callbacks(
+            AllocationChange.ALLOCATED, "/org/bluez/hci0/dev_44_44_33_11_23_12"
+        )
+
     assert len(ok_allocations) == 2
 
+    assert manager.async_current_allocations() == [
+        HaBluetoothSlotAllocations("AA:BB:CC:DD:EE:00", 5, 4, ["44:44:33:11:23:12"])
+    ]
+    assert manager.async_current_allocations("AA:BB:CC:DD:EE:00") == [
+        HaBluetoothSlotAllocations("AA:BB:CC:DD:EE:00", 5, 4, ["44:44:33:11:23:12"])
+    ]
     cancel1()
     cancel2()