From 0a9bef927c5f29c3e724fb60aa06706b6d896f82 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Tue, 21 Jan 2025 16:55:13 -1000 Subject: [PATCH] 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 | 2 ++ src/habluetooth/manager.pxd | 4 ++- src/habluetooth/manager.py | 62 +++++++++++++++++++++++++++++++------ src/habluetooth/models.py | 10 ++++++ tests/conftest.py | 14 +++++---- tests/test_manager.py | 50 +++++++++++++++++++++--------- 6 files changed, 111 insertions(+), 31 deletions(-) diff --git a/src/habluetooth/__init__.py b/src/habluetooth/__init__.py index da9d6fb..6c35773 100644 --- a/src/habluetooth/__init__.py +++ b/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", diff --git a/src/habluetooth/manager.pxd b/src/habluetooth/manager.pxd index a3ab647..23bfc1d 100644 --- a/src/habluetooth/manager.pxd +++ b/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( diff --git a/src/habluetooth/manager.py b/src/habluetooth/manager.py index ff465ff..0bf1bed 100644 --- a/src/habluetooth/manager.py +++ b/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] diff --git a/src/habluetooth/models.py b/src/habluetooth/models.py index 706eee9..ac7ca51 100644 --- a/src/habluetooth/models.py +++ b/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.""" diff --git a/tests/conftest.py b/tests/conftest.py index 269fd85..c4ee8ab 100644 --- a/tests/conftest.py +++ b/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,9 +160,9 @@ 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() @@ -168,8 +170,8 @@ def register_hci0_scanner() -> Generator[None, None, None]: @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() diff --git a/tests/test_manager.py b/tests/test_manager.py index 94caecc..51c5284 100644 --- a/tests/test_manager.py +++ b/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,9 +265,10 @@ 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"], @@ -270,25 +276,41 @@ def _ok_callback(allocations: Allocations) -> None: ) 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()