feat: Add ability to listen for ready devices (#664)

* feat: Add ability to listen for ready devices

Add a callback that will be invoked when devices become ready. This is to allow non-blocking setup with dynamically connected devices. In the future we can periodically or asynchronously invoke discover_devices and automatically find new devices as they become available.

* fix: Update roborock/devices/device.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* chore: add more tests for already connected devices

* chore: fix lint errors

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: allenporter

roborock/devices/device.py

@@ -11,6 +11,7 @@
 from collections.abc import Callable, Mapping
 from typing import Any, TypeVar, cast
 
+from roborock.callbacks import CallbackList
 from roborock.data import HomeDataDevice, HomeDataProduct
 from roborock.exceptions import RoborockException
 from roborock.roborock_message import RoborockMessage
@@ -23,6 +24,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 __all__ = [
+    "DeviceReadyCallback",
     "RoborockDevice",
 ]
 
@@ -33,6 +35,9 @@
 START_ATTEMPT_TIMEOUT = datetime.timedelta(seconds=5)
 
 
+DeviceReadyCallback = Callable[["RoborockDevice"], None]
+
+
 class RoborockDevice(ABC, TraitsMixin):
     """A generic channel for establishing a connection with a Roborock device.
 
@@ -67,6 +72,8 @@ def __init__(
         self._channel = channel
         self._connect_task: asyncio.Task[None] | None = None
         self._unsub: Callable[[], None] | None = None
+        self._ready_callbacks = CallbackList["RoborockDevice"]()
+        self._has_connected = False
 
     @property
     def duid(self) -> str:
@@ -110,6 +117,22 @@ def is_local_connected(self) -> bool:
         """
         return self._channel.is_local_connected
 
+    def add_ready_callback(self, callback: DeviceReadyCallback) -> Callable[[], None]:
+        """Add a callback to be notified when the device is ready.
+
+        A device is considered ready when it has successfully connected. It may go
+        offline later, but this callback will only be called once when the device
+        first connects.
+
+        The callback will be called immediately if the device has already previously
+        connected.
+        """
+        remove = self._ready_callbacks.add_callback(callback)
+        if self._has_connected:
+            callback(self)
+
+        return remove
+
     async def start_connect(self) -> None:
         """Start a background task to connect to the device.
 
@@ -133,6 +156,8 @@ async def connect_loop() -> None:
                     try:
                         await self.connect()
                         start_attempt.set()
+                        self._has_connected = True
+                        self._ready_callbacks(self)
                         return
                     except RoborockException as e:
                         start_attempt.set()

roborock/devices/device_manager.py

@@ -14,7 +14,7 @@
     HomeDataProduct,
     UserData,
 )
-from roborock.devices.device import RoborockDevice
+from roborock.devices.device import DeviceReadyCallback, RoborockDevice
 from roborock.map.map_parser import MapParserConfig
 from roborock.mqtt.roborock_session import create_lazy_mqtt_session
 from roborock.mqtt.session import MqttSession
@@ -155,6 +155,7 @@ async def create_device_manager(
     cache: Cache | None = None,
     map_parser_config: MapParserConfig | None = None,
     session: aiohttp.ClientSession | None = None,
+    ready_callback: DeviceReadyCallback | None = None,
 ) -> DeviceManager:
     """Convenience function to create and initialize a DeviceManager.
 
@@ -163,6 +164,7 @@ async def create_device_manager(
         cache: Optional cache implementation to use for caching device data.
         map_parser_config: Optional configuration for parsing maps.
         session: Optional aiohttp ClientSession to use for HTTP requests.
+        ready_callback: Optional callback to be notified when a device is ready.
 
     Returns:
         An initialized DeviceManager with discovered devices.
@@ -211,7 +213,11 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                     raise NotImplementedError(f"Device {device.name} has unsupported B01 model: {product.model}")
             case _:
                 raise NotImplementedError(f"Device {device.name} has unsupported version {device.pv}")
-        return RoborockDevice(device, product, channel, trait)
+
+        dev = RoborockDevice(device, product, channel, trait)
+        if ready_callback:
+            dev.add_ready_callback(ready_callback)
+        return dev
 
     manager = DeviceManager(web_api, device_creator, mqtt_session=mqtt_session, cache=cache)
     await manager.discover_devices()

tests/devices/test_device_manager.py

@@ -9,6 +9,7 @@
 
 from roborock.data import HomeData, UserData
 from roborock.devices.cache import InMemoryCache
+from roborock.devices.device import RoborockDevice
 from roborock.devices.device_manager import UserParams, create_device_manager, create_web_api_wrapper
 from roborock.exceptions import RoborockException
 
@@ -171,9 +172,30 @@ async def mock_home_data_with_counter(*args, **kwargs) -> HomeData:
         await device_manager.close()
 
 
+async def test_ready_callback(home_data: HomeData) -> None:
+    """Test that the ready callback is invoked when a device connects."""
+    ready_devices: list[RoborockDevice] = []
+    device_manager = await create_device_manager(USER_PARAMS, ready_callback=ready_devices.append)
+
+    # Callback should be called for the discovered device
+    assert len(ready_devices) == 1
+    device = ready_devices[0]
+    assert device.duid == "abc123"
+
+    # Verify that adding a ready callback to an already connected device will
+    # invoke the callback immediately.
+    more_ready_device: list[RoborockDevice] = []
+    device.add_ready_callback(more_ready_device.append)
+    assert len(more_ready_device) == 1
+    assert more_ready_device[0].duid == "abc123"
+
+    await device_manager.close()
+
+
 async def test_start_connect_failure(home_data: HomeData, channel_failure: Mock, mock_sleep: Mock) -> None:
     """Test that start_connect retries when connection fails."""
-    device_manager = await create_device_manager(USER_PARAMS)
+    ready_devices: list[RoborockDevice] = []
+    device_manager = await create_device_manager(USER_PARAMS, ready_callback=ready_devices.append)
     devices = await device_manager.get_devices()
 
     # The device should attempt to connect in the background at least once
@@ -184,6 +206,7 @@ async def test_start_connect_failure(home_data: HomeData, channel_failure: Mock,
     # Device should exist but not be connected
     assert len(devices) == 1
     assert not devices[0].is_connected
+    assert not ready_devices
 
     # Verify retry attempts
     assert channel_failure.return_value.subscribe.call_count >= 1
@@ -203,6 +226,8 @@ async def test_start_connect_failure(home_data: HomeData, channel_failure: Mock,
         assert attempts < 10, "Device did not connect after multiple attempts"
 
     assert devices[0].is_connected
+    assert ready_devices
+    assert len(ready_devices) == 1
 
     await device_manager.close()
     assert mock_unsub.call_count == 1