From 99fcb46a73ea6cb8f01817263d01a342365be78f Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Tue, 28 Jan 2025 08:43:11 -1000 Subject: [PATCH] feat: add async_register_scanner_registration_callback and async_current_scanners to the manager (#125) --- src/habluetooth/__init__.py | 6 ++++ src/habluetooth/base_scanner.pxd | 1 + src/habluetooth/base_scanner.py | 15 +++++--- src/habluetooth/manager.pxd | 1 + src/habluetooth/manager.py | 51 ++++++++++++++++++++++++++- src/habluetooth/models.py | 27 ++++++++++++++ tests/test_base_scanner.py | 14 +++++++- tests/test_manager.py | 60 +++++++++++++++++++++++++++++++- 8 files changed, 168 insertions(+), 7 deletions(-) diff --git a/src/habluetooth/__init__.py b/src/habluetooth/__init__.py index 2a37114..aea6537 100644 --- a/src/habluetooth/__init__.py +++ b/src/habluetooth/__init__.py @@ -19,6 +19,9 @@ BluetoothServiceInfoBleak, HaBluetoothConnector, HaBluetoothSlotAllocations, + HaScannerDetails, + HaScannerRegistration, + HaScannerRegistrationEvent, ) from .scanner import BluetoothScanningMode, HaScanner, ScannerStartError from .scanner_device import BluetoothScannerDevice @@ -44,6 +47,9 @@ "HaBluetoothConnector", "HaBluetoothSlotAllocations", "HaScanner", + "HaScannerDetails", + "HaScannerRegistration", + "HaScannerRegistrationEvent", "ScannerStartError", "get_manager", "set_manager", diff --git a/src/habluetooth/base_scanner.pxd b/src/habluetooth/base_scanner.pxd index 04a1cb1..a69d966 100644 --- a/src/habluetooth/base_scanner.pxd +++ b/src/habluetooth/base_scanner.pxd @@ -24,6 +24,7 @@ cdef class BaseHaScanner: cdef public object _cancel_watchdog cdef public object _loop cdef BluetoothManager _manager + cdef public object details cpdef tuple get_discovered_device_advertisement_data(self, str address) diff --git a/src/habluetooth/base_scanner.py b/src/habluetooth/base_scanner.py index 52b41c4..be5ced4 100644 --- a/src/habluetooth/base_scanner.py +++ b/src/habluetooth/base_scanner.py @@ -20,7 +20,7 @@ SCANNER_WATCHDOG_INTERVAL, SCANNER_WATCHDOG_TIMEOUT, ) -from .models import BluetoothServiceInfoBleak, HaBluetoothConnector +from .models import BluetoothServiceInfoBleak, HaBluetoothConnector, HaScannerDetails SCANNER_WATCHDOG_INTERVAL_SECONDS: Final = SCANNER_WATCHDOG_INTERVAL.total_seconds() _LOGGER = logging.getLogger(__name__) @@ -44,6 +44,7 @@ class BaseHaScanner: "adapter", "connectable", "connector", + "details", "name", "scanning", "source", @@ -54,9 +55,10 @@ def __init__( source: str, adapter: str, connector: HaBluetoothConnector | None = None, + connectable: bool = False, ) -> None: """Initialize the scanner.""" - self.connectable = False + self.connectable = connectable self.source = source self.connector = connector self._connecting = 0 @@ -68,6 +70,12 @@ def __init__( self._cancel_watchdog: asyncio.TimerHandle | None = None self._loop: asyncio.AbstractEventLoop | None = None self._manager = get_manager() + self.details = HaScannerDetails( + source=self.source, + connectable=self.connectable, + name=self.name, + adapter=self.adapter, + ) def async_setup(self) -> CALLBACK_TYPE: """Set up the scanner.""" @@ -211,8 +219,7 @@ def __init__( connectable: bool, ) -> None: """Initialize the scanner.""" - super().__init__(scanner_id, name, connector) - self.connectable = connectable + super().__init__(scanner_id, name, connector, connectable) self._details: dict[str, str | HaBluetoothConnector] = {"source": scanner_id} # Scanners only care about connectable devices. The manager # will handle taking care of availability for non-connectable devices diff --git a/src/habluetooth/manager.pxd b/src/habluetooth/manager.pxd index 23bfc1d..9af77b9 100644 --- a/src/habluetooth/manager.pxd +++ b/src/habluetooth/manager.pxd @@ -55,6 +55,7 @@ cdef class BluetoothManager: cdef public object _cancel_allocation_callbacks cdef public dict _adapter_sources cdef public dict _allocations + cdef public dict _scanner_registration_callbacks @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 f0d2c7f..0f83ef4 100644 --- a/src/habluetooth/manager.py +++ b/src/habluetooth/manager.py @@ -35,7 +35,12 @@ FALLBACK_MAXIMUM_STALE_ADVERTISEMENT_SECONDS, UNAVAILABLE_TRACK_SECONDS, ) -from .models import BluetoothServiceInfoBleak, HaBluetoothSlotAllocations +from .models import ( + BluetoothServiceInfoBleak, + HaBluetoothSlotAllocations, + HaScannerRegistration, + HaScannerRegistrationEvent, +) from .scanner_device import BluetoothScannerDevice from .usage import install_multiple_bleak_catcher, uninstall_multiple_bleak_catcher from .util import async_reset_adapter @@ -121,6 +126,7 @@ class BluetoothManager: "_loop", "_non_connectable_scanners", "_recovery_lock", + "_scanner_registration_callbacks", "_sources", "_unavailable_callbacks", "shutdown", @@ -171,6 +177,9 @@ def __init__( self._allocations_callbacks: dict[ str | None, set[Callable[[HaBluetoothSlotAllocations], None]] ] = {} + self._scanner_registration_callbacks: dict[ + str | None, set[Callable[[HaScannerRegistration], None]] + ] = {} @property def supports_passive_scan(self) -> bool: @@ -708,6 +717,7 @@ def _async_unregister_scanner_internal( self._allocations.pop(scanner.source, None) if connection_slots: self.slot_manager.remove_adapter(scanner.adapter) + self._async_on_scanner_registration(scanner, HaScannerRegistrationEvent.REMOVED) def async_register_scanner( self, @@ -728,6 +738,7 @@ def async_register_scanner( self._adapter_sources[scanner.adapter] = scanner.source if connection_slots: self.slot_manager.register_adapter(scanner.adapter, connection_slots) + self._async_on_scanner_registration(scanner, HaScannerRegistrationEvent.ADDED) return partial( self._async_unregister_scanner_internal, scanners, scanner, connection_slots ) @@ -797,6 +808,23 @@ def async_on_allocation_changed(self, allocations: Allocations) -> None: except Exception: _LOGGER.exception("Error in allocation callback") + def _async_on_scanner_registration( + self, scanner: BaseHaScanner, event: HaScannerRegistrationEvent + ) -> None: + """Call scanner callbacks.""" + for source_key in (scanner.source, None): + if not ( + scanner_callbacks := self._scanner_registration_callbacks.get( + source_key + ) + ): + continue + for callback_ in scanner_callbacks: + try: + callback_(HaScannerRegistration(event, scanner)) + except Exception: + _LOGGER.exception("Error in scanner callback") + def async_current_allocations( self, source: str | None = None ) -> list[HaBluetoothSlotAllocations] | None: @@ -823,3 +851,24 @@ def _async_unregister_allocation_callback( callbacks.discard(callback) if not callbacks: del self._allocations_callbacks[source] + + def async_register_scanner_registration_callback( + self, callback: Callable[[HaScannerRegistration], None], source: str | None + ) -> CALLBACK_TYPE: + """Register a callback to be called when a scanner is added or removed.""" + self._scanner_registration_callbacks.setdefault(source, set()).add(callback) + return partial( + self._async_unregister_scanner_registration_callback, callback, source + ) + + def _async_unregister_scanner_registration_callback( + self, callback: Callable[[HaScannerRegistration], None], source: str | None + ) -> None: + if (callbacks := self._scanner_registration_callbacks.get(source)) is not None: + callbacks.discard(callback) + if not callbacks: + del self._scanner_registration_callbacks[source] + + def async_current_scanners(self) -> list[BaseHaScanner]: + """Return the current scanners.""" + return list(self._sources.values()) diff --git a/src/habluetooth/models.py b/src/habluetooth/models.py index df8cc4a..5906330 100644 --- a/src/habluetooth/models.py +++ b/src/habluetooth/models.py @@ -13,6 +13,7 @@ from bleak_retry_connector import NO_RSSI_VALUE if TYPE_CHECKING: + from .base_scanner import BaseHaScanner from .manager import BluetoothManager _BluetoothServiceInfoSelfT = TypeVar( @@ -57,6 +58,22 @@ class HaBluetoothSlotAllocations: allocated: list[str] # Addresses of connected devices +class HaScannerRegistrationEvent(Enum): + """Events for scanner registration.""" + + ADDED = "added" + REMOVED = "removed" + UPDATED = "updated" + + +@dataclass(slots=True, frozen=True) +class HaScannerRegistration: + """Data for a scanner event.""" + + event: HaScannerRegistrationEvent + scanner: BaseHaScanner + + @dataclass(slots=True) class HaBluetoothConnector: """Data for how to connect a BLEDevice from a given scanner.""" @@ -66,6 +83,16 @@ class HaBluetoothConnector: can_connect: Callable[[], bool] +@dataclass(slots=True, frozen=True) +class HaScannerDetails: + """Details for a scanner.""" + + source: str + connectable: bool + name: str + adapter: str + + class BluetoothScanningMode(Enum): """The mode of scanning for bluetooth devices.""" diff --git a/tests/test_base_scanner.py b/tests/test_base_scanner.py index 0a273d0..7bd4f0b 100644 --- a/tests/test_base_scanner.py +++ b/tests/test_base_scanner.py @@ -11,7 +11,12 @@ from bleak.backends.scanner import AdvertisementData from bluetooth_data_tools import monotonic_time_coarse -from habluetooth import BaseHaRemoteScanner, HaBluetoothConnector, get_manager +from habluetooth import ( + BaseHaRemoteScanner, + HaBluetoothConnector, + HaScannerDetails, + get_manager, +) from habluetooth.const import ( CONNECTABLE_FALLBACK_MAXIMUM_STALE_ADVERTISEMENT_SECONDS, FALLBACK_MAXIMUM_STALE_ADVERTISEMENT_SECONDS, @@ -103,6 +108,13 @@ async def test_remote_scanner(name_2: str | None) -> None: MockBleakClient, "mock_bleak_client", lambda: False ) scanner = FakeScanner("esp32", "esp32", connector, True) + details = scanner.details + assert details == HaScannerDetails( + source=scanner.source, + connectable=scanner.connectable, + name=scanner.name, + adapter=scanner.adapter, + ) unsetup = scanner.async_setup() cancel = manager.async_register_scanner(scanner) diff --git a/tests/test_manager.py b/tests/test_manager.py index 563249a..b8ba3c1 100644 --- a/tests/test_manager.py +++ b/tests/test_manager.py @@ -13,6 +13,8 @@ from habluetooth import ( BluetoothManager, HaBluetoothSlotAllocations, + HaScannerRegistration, + HaScannerRegistrationEvent, get_manager, set_manager, ) @@ -24,7 +26,7 @@ inject_advertisement_with_source, utcnow, ) -from .conftest import FakeBluetoothAdapters +from .conftest import FakeBluetoothAdapters, FakeScanner @pytest.mark.asyncio @@ -332,3 +334,59 @@ async def test_async_register_allocation_callback_non_connectable( allocated=[], ), ] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("enable_bluetooth") +async def test_async_register_scanner_registration_callback( + register_hci0_scanner: None, + register_hci1_scanner: None, +) -> None: + """Test bluetooth async_register_scanner_registration_callback handles failures.""" + manager = get_manager() + assert manager._loop is not None + + scanners = manager.async_current_scanners() + assert len(scanners) == 2 + sources = {scanner.source for scanner in scanners} + assert sources == {"AA:BB:CC:DD:EE:00", "AA:BB:CC:DD:EE:11"} + + failed_scanner_callbacks: list[HaScannerRegistration] = [] + + def _failing_callback(scanner_registration: HaScannerRegistration) -> None: + """Failing callback.""" + failed_scanner_callbacks.append(scanner_registration) + raise ValueError("This is a test") + + ok_scanner_callbacks: list[HaScannerRegistration] = [] + + def _ok_callback(scanner_registration: HaScannerRegistration) -> None: + """Ok callback.""" + ok_scanner_callbacks.append(scanner_registration) + + cancel1 = manager.async_register_scanner_registration_callback( + _failing_callback, None + ) + # Make sure the second callback still works if the first one fails and + # raises an exception + cancel2 = manager.async_register_scanner_registration_callback(_ok_callback, None) + + hci3_scanner = FakeScanner("AA:BB:CC:DD:EE:33", "hci3") + hci3_scanner.connectable = True + manager = get_manager() + cancel = manager.async_register_scanner(hci3_scanner, connection_slots=5) + + assert len(ok_scanner_callbacks) == 1 + assert ok_scanner_callbacks[0] == HaScannerRegistration( + HaScannerRegistrationEvent.ADDED, hci3_scanner + ) + assert len(failed_scanner_callbacks) == 1 + + cancel() + + assert len(ok_scanner_callbacks) == 2 + assert ok_scanner_callbacks[1] == HaScannerRegistration( + HaScannerRegistrationEvent.REMOVED, hci3_scanner + ) + cancel1() + cancel2()