feat: add Q10 integration test and support stream streaming on FakeChannel

Author: allenporter

roborock/testing/channel.py

@@ -5,7 +5,8 @@
 in-memory replacement for `MqttChannel` and `LocalChannel` during testing.
 """
 
-from collections.abc import Awaitable, Callable
+import asyncio
+from collections.abc import AsyncGenerator, Awaitable, Callable
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock
 
@@ -146,3 +147,17 @@ def clear_error(self) -> None:
         self.publish.side_effect = self._publish
         self.subscribe.side_effect = self._subscribe
         self.connect.side_effect = self._connect
+
+    async def subscribe_stream(self) -> AsyncGenerator[RoborockMessage, None]:
+        """Stream messages received via this channel."""
+        queue: asyncio.Queue[RoborockMessage] = asyncio.Queue()
+
+        def callback(message: RoborockMessage) -> None:
+            queue.put_nowait(message)
+
+        unsub = await self.subscribe(callback)
+        try:
+            while True:
+                yield await queue.get()
+        finally:
+            unsub()

roborock/testing/cloud.py

@@ -16,6 +16,7 @@
 from roborock.data import HomeData, Reference, RRiot, UserData
 from roborock.devices.rpc.v1_channel import create_v1_channel as original_create_v1_channel
 from roborock.devices.transport.mqtt_channel import create_mqtt_channel as original_create_mqtt_channel
+from roborock.testing.b01_q10_simulator import Q10VacuumSimulator
 from roborock.testing.simulator import RoborockDeviceSimulator
 from roborock.testing.v1_simulator import V1VacuumSimulator
 
@@ -261,6 +262,12 @@ def mock_create_v1_channel(user_data, mqtt_params, mqtt_session, device, device_
         def mock_create_mqtt_channel(user_data, mqtt_params, mqtt_session, device):
             server = self.simulated_devices.get(device.duid)
             if server:
+                if device.pv == "B01" and not isinstance(server, Q10VacuumSimulator):
+                    raise NotImplementedError(
+                        f"Simulating protocol {device.pv} is not yet supported. "
+                        "Only Q10VacuumSimulator is supported for B01 protocols currently."
+                    )
+                server.mqtt_channel._is_connected = True
                 return server.mqtt_channel
             if device.pv in ("A01", "B01"):
                 raise NotImplementedError(

tests/devices/test_device_manager.py

@@ -15,7 +15,7 @@
 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, RoborockInvalidCredentials
-from roborock.testing import FakeRoborockCloud, V1VacuumSimulator
+from roborock.testing import FakeRoborockCloud, Q10VacuumSimulator, V1VacuumSimulator
 from tests import mock_data
 
 USER_DATA = UserData.from_dict(mock_data.USER_DATA)
@@ -56,6 +56,83 @@ async def test_no_devices(cloud: FakeRoborockCloud, patch_device_manager: None)
     assert devices == []
 
 
+async def test_with_q10_device(cloud: FakeRoborockCloud, patch_device_manager: None) -> None:
+    """Test DeviceManager with a Q10 device simulator registered."""
+    product = HomeDataProduct(
+        id="product-id-q10",
+        name="Roborock Q10",
+        model="roborock.vacuum.ss07",
+        category=RoborockCategory.VACUUM,
+    )
+    device_info = HomeDataDevice(
+        duid="q10_duid",
+        name="My Q10",
+        local_key="key123key123key1",
+        product_id=product.id,
+        sn="q10_serial",
+        pv="B01",
+    )
+
+    home_data = cloud.web_api.get_default_home_data()
+    home_data.devices.append(device_info)
+    home_data.products.append(product)
+    cloud.web_api.set_homes_response(home_data)
+
+    q10_sim = Q10VacuumSimulator(
+        duid="q10_duid",
+        device_info=device_info,
+        product=product,
+    )
+    cloud.add_device(q10_sim)
+
+    device_manager = await create_device_manager(USER_PARAMS)
+    devices = await device_manager.get_devices()
+
+    # The setup includes fake_device (V1) by default because of the fake_device fixture
+    # which is not requested here (we only request cloud and patch_device_manager),
+    # but the mock EAPI returns the full home_data layout containing our Q10 device.
+    assert len(devices) == 1
+
+    q10_device = await device_manager.get_device("q10_duid")
+    assert q10_device is not None
+    assert q10_device.name == "My Q10"
+    assert q10_device.product.model == "roborock.vacuum.ss07"
+
+    # Wait for background connect to establish
+    for _ in range(20):
+        if q10_device.is_connected:
+            break
+        await asyncio.sleep(0.05)
+    assert q10_device.is_connected
+
+    assert q10_device.b01_q10_properties is not None
+    assert q10_device.b01_q10_properties.status.status is None
+
+    await q10_device.b01_q10_properties.refresh()
+
+    for _ in range(20):
+        if q10_device.b01_q10_properties.status.battery == 100:
+            break
+        await asyncio.sleep(0.05)
+
+    assert q10_device.b01_q10_properties.status.battery == 100
+    from roborock.data.b01_q10.b01_q10_code_mappings import YXDeviceState
+
+    assert q10_device.b01_q10_properties.status.status == YXDeviceState.CHARGING
+
+    await q10_device.b01_q10_properties.vacuum.start_clean()
+
+    for _ in range(20):
+        if q10_device.b01_q10_properties.status.status == YXDeviceState.CLEANING:
+            break
+        await asyncio.sleep(0.05)
+
+    assert q10_device.b01_q10_properties.status.status == YXDeviceState.CLEANING
+    assert q10_sim.status[121] == 5
+
+    await device_manager.close()
+
+
 async def test_with_device(
     cloud: FakeRoborockCloud, fake_device: V1VacuumSimulator, patch_device_manager: None
 ) -> None: