feat: Recognize Q10 devices and add a command trait

Add the ability to send commands to roborock CLI. This adds a single trait for sending commands, using a blend of approaches from #692 and #709

Author: allenporter

roborock/cli.py

@@ -43,6 +43,7 @@
 
 from roborock import RoborockCommand
 from roborock.data import RoborockBase, UserData
+from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
 from roborock.device_features import DeviceFeatures
 from roborock.devices.cache import Cache, CacheData
 from roborock.devices.device import RoborockDevice
@@ -745,6 +746,21 @@ async def network_info(ctx, device_id: str):
     await _display_v1_trait(context, device_id, lambda v1: v1.network_info)
 
 
+def _parse_b01_q10_command(cmd: str) -> B01_Q10_DP:
+    """Parse B01_Q10 command from either enum name or value."""
+    try:
+        return B01_Q10_DP(int(cmd))
+    except ValueError:
+        try:
+            return B01_Q10_DP.from_name(cmd)
+        except ValueError:
+            try:
+                return B01_Q10_DP.from_value(cmd)
+            except ValueError:
+                pass
+    raise RoborockException(f"Invalid command {cmd} for B01_Q10 device")
+
+
 @click.command()
 @click.option("--device_id", required=True)
 @click.option("--cmd", required=True)
@@ -755,12 +771,18 @@ async def command(ctx, cmd, device_id, params):
     context: RoborockContext = ctx.obj
     device_manager = await context.get_device_manager()
     device = await device_manager.get_device(device_id)
-    if device.v1_properties is None:
-        raise RoborockException(f"Device {device.name} does not support V1 protocol")
-    command_trait: Trait = device.v1_properties.command
-    result = await command_trait.send(cmd, json.loads(params) if params is not None else None)
-    if result:
-        click.echo(dump_json(result))
+    if device.v1_properties is not None:
+        command_trait: Trait = device.v1_properties.command
+        result = await command_trait.send(cmd, json.loads(params) if params is not None else None)
+        if result:
+            click.echo(dump_json(result))
+    elif device.b01_q10_properties is not None:
+        cmd_value = _parse_b01_q10_command(cmd)
+        command_trait: Trait = device.b01_q10_properties.command
+        await command_trait.send(cmd_value, json.loads(params) if params is not None else None)
+        click.echo("Command sent successfully; Enable debug logging (-d) to see responses.")
+        # Q10 commands don't have a specific time to respond, so wait a bit and log
+        await asyncio.sleep(5)
 
 
 @click.command()

roborock/devices/b01_q10_channel.py

@@ -0,0 +1,38 @@
+"""Thin wrapper around the MQTT channel for Roborock B01 Q10 devices."""
+
+from __future__ import annotations
+
+import logging
+
+from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
+from roborock.exceptions import RoborockException
+from roborock.protocols.b01_q10_protocol import (
+    ParamsType,
+    encode_mqtt_payload,
+)
+
+from .mqtt_channel import MqttChannel
+
+_LOGGER = logging.getLogger(__name__)
+_TIMEOUT = 10.0
+
+
+async def send_command(
+    mqtt_channel: MqttChannel,
+    command: B01_Q10_DP,
+    params: ParamsType,
+) -> None:
+    """Send a command on the MQTT channel, without waiting for a response"""
+    _LOGGER.debug("Sending B01 MQTT command: cmd=%s params=%s", command, params)
+    roborock_message = encode_mqtt_payload(command, params)
+    _LOGGER.debug("Sending MQTT message: %s", roborock_message)
+    try:
+        await mqtt_channel.publish(roborock_message)
+    except RoborockException as ex:
+        _LOGGER.debug(
+            "Error sending B01 decoded command (method=%s params=%s): %s",
+            command,
+            params,
+            ex,
+        )
+        raise

roborock/devices/b01_q7_channel.py

@@ -1,4 +1,4 @@
-"""Thin wrapper around the MQTT channel for Roborock B01 devices."""
+"""Thin wrapper around the MQTT channel for Roborock B01 Q7 devices."""
 
 from __future__ import annotations
 

roborock/devices/device_manager.py

@@ -242,9 +242,7 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                 channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
                 model_part = product.model.split(".")[-1]
                 if "ss" in model_part:
-                    raise UnsupportedDeviceError(
-                        f"Device {device.name} has unsupported version B01 product model {product.model}"
-                    )
+                    trait = b01.q10.create(channel)
                 elif "sc" in model_part:
                     # Q7 devices start with 'sc' in their model naming.
                     trait = b01.q7.create(channel)

roborock/devices/traits/b01/__init__.py

@@ -1,5 +1,12 @@
 """Traits for B01 devices."""
 
+from . import q7, q10
 from .q7 import Q7PropertiesApi
+from .q10 import Q10PropertiesApi
 
-__all__ = ["Q7PropertiesApi", "q7", "q10"]
+__all__ = [
+    "Q7PropertiesApi",
+    "Q10PropertiesApi",
+    "q7",
+    "q10",
+]

roborock/devices/traits/b01/q10/__init__.py

@@ -1 +1,29 @@
-"""Q10"""
+"""Traits for Q10 B01 devices."""
+
+from typing import Any
+
+from roborock.devices.b01_q7_channel import send_decoded_command
+from roborock.devices.mqtt_channel import MqttChannel
+from roborock.devices.traits import Trait
+
+from .command import CommandTrait
+
+__all__ = [
+    "Q10PropertiesApi",
+]
+
+
+class Q10PropertiesApi(Trait):
+    """API for interacting with B01 devices."""
+
+    command: CommandTrait
+    """Trait for sending commands to Q10 devices."""
+
+    def __init__(self, channel: MqttChannel) -> None:
+        """Initialize the B01Props API."""
+        self.command = CommandTrait(channel)
+
+
+def create(channel: MqttChannel) -> Q10PropertiesApi:
+    """Create traits for B01 devices."""
+    return Q10PropertiesApi(channel)

roborock/devices/traits/b01/q10/command.py

@@ -0,0 +1,32 @@
+from typing import Any
+
+from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
+from roborock.devices.b01_q10_channel import send_command
+from roborock.devices.mqtt_channel import MqttChannel
+from roborock.protocols.b01_q10_protocol import ParamsType
+
+
+class CommandTrait:
+    """Trait for sending commands to Q10 Roborock devices.
+
+    This trait allows sending raw commands directly to the device. It is particularly
+    useful for accessing features that do not have their own traits. Generally
+    it is preferred to use specific traits for device functionality when
+    available.
+    """
+
+    def __init__(self, channel: MqttChannel) -> None:
+        """Initialize the CommandTrait."""
+        self._channel = channel
+
+    async def send(self, command: B01_Q10_DP, params: ParamsType = None) -> Any:
+        """Send a command to the device.
+
+        Sending a raw command to the device using this method does not update
+        the internal state of any other traits. It is the responsibility of the
+        caller to ensure that any traits affected by the command are refreshed
+        as needed.
+        """
+        if not self._channel:
+            raise ValueError("Device trait in invalid state")
+        return await send_command(self._channel, command, params=params)

roborock/devices/traits/traits_mixin.py

@@ -34,6 +34,9 @@ class TraitsMixin:
     b01_q7_properties: b01.Q7PropertiesApi | None = None
     """B01 Q7 properties trait, if supported."""
 
+    b01_q10_properties: b01.Q10PropertiesApi | None = None
+    """B01 Q10 properties trait, if supported."""
+
     def __init__(self, trait: Trait) -> None:
         """Initialize the TraitsMixin with the given trait.