feat: Add a trait for sending commands (#539)

* feat: Add a trait for sending commands

* feat: Update cli to use new command interface

* chore: fix tests

* Update tests/devices/traits/v1/test_command.py

Co-authored-by: Luke Lashley <conway220@gmail.com>

* Update tests/devices/traits/v1/test_command.py

Co-authored-by: Luke Lashley <conway220@gmail.com>

* chore: fix syntax and lint errors

* chore: fix lint errors

---------

Co-authored-by: Luke Lashley <conway220@gmail.com>

Author: allenporter

roborock/cli.py

@@ -714,21 +714,14 @@ async def home(ctx, device_id: str, refresh: bool):
 @async_command
 async def command(ctx, cmd, device_id, params):
     context: RoborockContext = ctx.obj
-    cache_data = await context.get_devices()
-
-    home_data = cache_data.home_data
-    devices = home_data.get_all_devices()
-    device = next(device for device in devices if device.duid == device_id)
-    model = next(
-        (product.model for product in home_data.products if device is not None and product.id == device.product_id),
-        None,
-    )
-    if model is None:
-        raise RoborockException(f"Could not find model for device {device.name}")
-    device_info = DeviceData(device=device, model=model)
-    mqtt_client = RoborockMqttClientV1(cache_data.user_data, device_info)
-    await mqtt_client.send_command(cmd, json.loads(params) if params is not None else None)
-    await mqtt_client.async_release()
+    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))
 
 
 @click.command()

roborock/devices/traits/v1/__init__.py

@@ -12,6 +12,7 @@
 
 from .child_lock import ChildLockTrait
 from .clean_summary import CleanSummaryTrait
+from .command import CommandTrait
 from .common import V1TraitMixin
 from .consumeable import ConsumableTrait
 from .device_features import DeviceFeaturesTrait
@@ -39,6 +40,7 @@
     "ConsumableTrait",
     "HomeTrait",
     "DeviceFeaturesTrait",
+    "CommandTrait",
     "ChildLockTrait",
     "FlowLedStatusTrait",
     "LedStatusTrait",
@@ -54,6 +56,7 @@ class PropertiesApi(Trait):
 
     # All v1 devices have these traits
     status: StatusTrait
+    command: CommandTrait
     dnd: DoNotDisturbTrait
     clean_summary: CleanSummaryTrait
     sound_volume: SoundVolumeTrait

roborock/devices/traits/v1/command.py

@@ -0,0 +1,21 @@
+from typing import Any
+
+from roborock import RoborockCommand
+
+
+class CommandTrait:
+    """Trait for sending commands to Roborock devices."""
+
+    def __post_init__(self) -> None:
+        """Post-initialization to set up the RPC channel.
+
+        This is called automatically after the dataclass is initialized by the
+        device setup code.
+        """
+        self._rpc_channel = None
+
+    async def send(self, command: RoborockCommand, params: dict[str, Any] | None = None) -> Any:
+        """Send a command to the device."""
+        if not self._rpc_channel:
+            raise ValueError("Device trait in invalid state")
+        return await self._rpc_channel.send_command(command, params=params)

tests/devices/traits/v1/test_command.py

@@ -0,0 +1,61 @@
+"""Tests for the CommandTrait class."""
+
+from unittest.mock import AsyncMock
+
+import pytest
+
+from roborock.devices.traits.v1.command import CommandTrait
+from roborock.exceptions import RoborockException
+from roborock.roborock_typing import RoborockCommand
+
+
+@pytest.fixture(name="command_trait")
+def command_trait_fixture() -> CommandTrait:
+    """Create a CommandTrait instance with a mocked RPC channel."""
+    trait = CommandTrait()
+    trait._rpc_channel = AsyncMock()  # type: ignore[assignment]
+    return trait
+
+
+async def test_send_command_success(command_trait: CommandTrait) -> None:
+    """Test successfully sending a command."""
+    mock_rpc_channel = command_trait._rpc_channel
+    assert mock_rpc_channel is not None
+    mock_rpc_channel.send_command.return_value = {"result": "ok"}
+
+    # Call the method
+    result = await command_trait.send(RoborockCommand.APP_START)
+
+    # Verify the result
+    assert result == {"result": "ok"}
+
+    # Verify the RPC call was made correctly
+    mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.APP_START, params=None)
+
+
+async def test_send_command_with_params(command_trait: CommandTrait) -> None:
+    """Test successfully sending a command with parameters."""
+    mock_rpc_channel = command_trait._rpc_channel
+    assert mock_rpc_channel is not None
+    mock_rpc_channel.send_command.return_value = {"result": "ok"}
+    params = {"segments": [1, 2, 3]}
+
+    # Call the method
+    result = await command_trait.send(RoborockCommand.APP_SEGMENT_CLEAN, params)
+
+    # Verify the result
+    assert result == {"result": "ok"}
+
+    # Verify the RPC call was made correctly
+    mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.APP_SEGMENT_CLEAN, params=params)
+
+
+async def test_send_command_propagates_exception(command_trait: CommandTrait) -> None:
+    """Test that exceptions from RPC channel are propagated."""
+    mock_rpc_channel = command_trait._rpc_channel
+    assert mock_rpc_channel is not None
+    mock_rpc_channel.send_command.side_effect = RoborockException("Communication error")
+
+    # Verify the exception is propagated
+    with pytest.raises(RoborockException, match="Communication error"):
+        await command_trait.send(RoborockCommand.APP_START)