chore: Only allow a single trait

Author: allenporter

roborock/devices/device.py

@@ -38,15 +38,15 @@ def __init__(
         self,
         device_info: HomeDataDevice,
         channel: Channel,
-        traits: list[Trait],
+        trait: Trait,
     ) -> None:
         """Initialize the RoborockDevice.
 
         The device takes ownership of the channel for communication with the device.
         Use `connect()` to establish the connection, which will set up the appropriate
         protocol channel. Use `close()` to clean up all connections.
         """
-        TraitsMixin.__init__(self, traits)
+        TraitsMixin.__init__(self, trait)
         self._duid = device_info.duid
         self._name = device_info.name
         self._channel = channel

roborock/devices/device_manager.py

@@ -145,23 +145,23 @@ async def create_device_manager(
 
     def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> RoborockDevice:
         channel: Channel
-        traits: list[Trait] = []
+        trait: Trait
         match device.pv:
             case DeviceVersion.V1:
                 v1_channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache)
                 channel = v1_channel
-                traits.extend(v1.create_v1_traits(product, v1_channel.rpc_channel))
+                trait = v1.create(product, v1_channel.rpc_channel)
             case DeviceVersion.A01:
                 mqtt_channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
                 channel = mqtt_channel
-                traits.extend(a01.create_a01_traits(product, mqtt_channel))
+                trait = a01.create(product, mqtt_channel)
             case DeviceVersion.B01:
                 mqtt_channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
                 channel = mqtt_channel
-                traits.extend(b01.create_b01_traits(mqtt_channel))
+                trait = b01.create(mqtt_channel)
             case _:
                 raise NotImplementedError(f"Device {device.name} has unsupported version {device.pv}")
-        return RoborockDevice(device, channel, traits)
+        return RoborockDevice(device, channel, trait)
 
     manager = DeviceManager(home_data_api, device_creator, mqtt_session=mqtt_session, cache=cache)
     await manager.discover_devices()

roborock/devices/traits/a01/__init__.py

@@ -50,12 +50,12 @@ async def set_value(self, protocol: RoborockZeoProtocol, value: Any) -> dict[Rob
         return await send_decoded_command(self._channel, params)
 
 
-def create_a01_traits(product: HomeDataProduct, mqtt_channel: MqttChannel) -> list[Trait]:
+def create(product: HomeDataProduct, mqtt_channel: MqttChannel) -> DyadApi | ZeoApi:
     """Create traits for A01 devices."""
     match product.category:
         case RoborockCategory.WET_DRY_VAC:
-            return [DyadApi(mqtt_channel)]
+            return DyadApi(mqtt_channel)
         case RoborockCategory.WASHING_MACHINE:
-            return [ZeoApi(mqtt_channel)]
+            return ZeoApi(mqtt_channel)
         case _:
             raise NotImplementedError(f"Unsupported category {product.category}")

roborock/devices/traits/b01/__init__.py

@@ -1,17 +1,31 @@
 """Traits for B01 devices."""
 
+from roborock import RoborockB01Methods
 from roborock.devices.b01_channel import send_decoded_command
 from roborock.devices.mqtt_channel import MqttChannel
 from roborock.devices.traits import Trait
-
-from .props import B01PropsApi
+from roborock.roborock_message import RoborockB01Props
 
 __init__ = [
     "create_b01_traits",
     "B01PropsApi",
 ]
 
 
-def create_b01_traits(channel: MqttChannel) -> list[Trait]:
+class B01PropsApi(Trait):
+    """API for interacting with B01 devices."""
+
+    def __init__(self, channel: MqttChannel) -> None:
+        """Initialize the B01Props API."""
+        self._channel = channel
+
+    async def query_values(self, props: list[RoborockB01Props]) -> None:
+        """Query the device for the values of the given Dyad protocols."""
+        await send_decoded_command(
+            self._channel, dps=10000, command=RoborockB01Methods.GET_PROP, params={"property": props}
+        )
+
+
+def create(channel: MqttChannel) -> B01PropsApi:
     """Create traits for B01 devices."""
-    return [B01PropsApi(channel)]
+    return B01PropsApi(channel)

roborock/devices/traits/b01/props.py

@@ -22,7 +22,7 @@
 class TraitsMixin:
     """Mixin to provide trait accessors."""
 
-    v1_properties: v1.Properties | None = None
+    v1_properties: v1.PropertiesApi | None = None
     """V1 properties trait, if supported."""
 
     dyad: a01.DyadApi | None = None
@@ -31,20 +31,21 @@ class TraitsMixin:
     zeo: a01.ZeoApi | None = None
     """Zeo API, if supported."""
 
-    b01_properties: b01.B01PropsApi | None = None
+    b01_props_api: b01.B01PropsApi | None = None
     """B01 properties trait, if supported."""
 
-    def __init__(self, traits: list[Trait]) -> None:
-        """Initialize the TraitsMixin with the given traits list.
+    def __init__(self, trait: Trait) -> None:
+        """Initialize the TraitsMixin with the given trait.
 
         This will populate the appropriate trait attributes based on the types
         of the traits provided.
         """
-        trait_map: dict[type[Trait], Trait] = {type(item): item for item in traits}
+        # trait_map: dict[type[Trait], Trait] = {type(item): item for item in traits}
         for item in fields(self):
             trait_type = _get_trait_type(item)
-            if (trait := trait_map.get(trait_type, None)) is not None:
+            if trait_type == type(trait):
                 setattr(self, item.name, trait)
+                break
 
 
 def _get_trait_type(item) -> type[Trait]:

roborock/devices/traits/traits_mixin.py

@@ -9,14 +9,14 @@
 from .properties import CleanSummaryTrait, DoNotDisturbTrait, SoundVolumeTrait, StatusTrait
 
 __all__ = [
-    "create_v1_traits",
-    "Properties",
+    "create",
+    "PropertiesApi",
     "properties",
 ]
 
 
 @dataclass
-class Properties(Trait):
+class PropertiesApi(Trait):
     """Common properties for V1 devices.
 
     This class holds all the traits that are common across all V1 devices.
@@ -34,16 +34,14 @@ def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel) -> None:
         """Initialize the V1TraitProps with None values."""
         self.status = StatusTrait(product)
 
+        # Create traits and set the RPC channel
         for item in fields(self):
             if (trait := getattr(self, item.name, None)) is None:
                 trait = item.type()
                 setattr(self, item.name, trait)
             trait._rpc_channel = rpc_channel
 
 
-def create_v1_traits(product: HomeDataProduct, rpc_channel: V1RpcChannel) -> list[Trait]:
+def create(product: HomeDataProduct, rpc_channel: V1RpcChannel) -> PropertiesApi:
     """Create traits for V1 devices."""
-    return [
-        Properties(product, rpc_channel)
-        # Add optional traits here as needed in the future
-    ]
+    return PropertiesApi(product, rpc_channel)

roborock/devices/traits/v1/__init__.py

@@ -9,8 +9,7 @@
 
 from roborock.containers import HomeData, S7MaxVStatus, UserData
 from roborock.devices.device import RoborockDevice
-from roborock.devices.traits import Trait
-from roborock.devices.traits.v1 import create_v1_traits
+from roborock.devices.traits import v1
 from roborock.devices.traits.v1.common import V1TraitMixin
 from roborock.devices.v1_rpc_channel import decode_rpc_response
 from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
@@ -37,21 +36,15 @@ def rpc_channel_fixture() -> AsyncMock:
 
 
 @pytest.fixture(autouse=True, name="device")
-def device_fixture(channel: AsyncMock, traits: list[Trait]) -> RoborockDevice:
+def device_fixture(channel: AsyncMock, rpc_channel: AsyncMock) -> RoborockDevice:
     """Fixture to set up the device for tests."""
     return RoborockDevice(
         device_info=HOME_DATA.devices[0],
         channel=channel,
-        traits=traits,
+        trait=v1.create(HOME_DATA.products[0], rpc_channel),
     )
 
 
-@pytest.fixture(autouse=True, name="traits")
-def traits_fixture(rpc_channel: AsyncMock) -> list[Trait]:
-    """Fixture to set up the V1 API for tests."""
-    return create_v1_traits(HOME_DATA.products[0], rpc_channel)
-
-
 async def test_device_connection(device: RoborockDevice, channel: AsyncMock) -> None:
     """Test the Device connection setup."""
 
@@ -91,7 +84,7 @@ def setup_rpc_channel_fixture(rpc_channel: AsyncMock, payload: pathlib.Path) ->
 
 
 @pytest.mark.parametrize(
-    ("payload", "trait_method"),
+    ("payload", "property_method"),
     [
         (TESTDATA / "get_status.json", lambda x: x.status),
         (TESTDATA / "get_dnd.json", lambda x: x.dnd),
@@ -103,11 +96,11 @@ async def test_device_trait_command_parsing(
     device: RoborockDevice,
     setup_rpc_channel: AsyncMock,
     snapshot: SnapshotAssertion,
-    trait_method: Callable[..., V1TraitMixin],
+    property_method: Callable[..., V1TraitMixin],
     payload: str,
 ) -> None:
     """Test the device trait command."""
-    trait = trait_method(device.v1_properties)
+    trait = property_method(device.v1_properties)
     assert trait
     assert isinstance(trait, V1TraitMixin)
     await trait.refresh()

tests/devices/test_v1_device.py

@@ -6,8 +6,7 @@
 
 from roborock.containers import HomeData, S7MaxVStatus, UserData
 from roborock.devices.device import RoborockDevice
-from roborock.devices.traits import Trait
-from roborock.devices.traits.v1 import create_v1_traits
+from roborock.devices.traits import v1
 
 from .... import mock_data
 
@@ -29,16 +28,10 @@ def rpc_channel_fixture() -> AsyncMock:
 
 
 @pytest.fixture(autouse=True, name="device")
-def device_fixture(channel: AsyncMock, traits: list[Trait]) -> RoborockDevice:
+def device_fixture(channel: AsyncMock, mock_rpc_channel: AsyncMock) -> RoborockDevice:
     """Fixture to set up the device for tests."""
     return RoborockDevice(
         device_info=HOME_DATA.devices[0],
         channel=channel,
-        traits=traits,
+        trait=v1.create(HOME_DATA.products[0], mock_rpc_channel),
     )
-
-
-@pytest.fixture(autouse=True, name="traits")
-def traits_fixture(mock_rpc_channel: AsyncMock) -> list[Trait]:
-    """Fixture to set up the V1 API for tests."""
-    return create_v1_traits(HOME_DATA.products[0], mock_rpc_channel)