chore: refactor to separate b01 q7 and q10 logic (#635)

* chore: refactor to seperate b01 ss and sc logic

* chore: refactor to seperate b01 ss and sc logic

* fix: clean up some naming

* fix: update roborock/devices/traits/b01/__init__.py

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

* chore: update roborock/devices/traits/b01/q7/__init__.py

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

* chore: update roborock/devices/device_manager.py

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

* fix: use strip not split

* chore: share duplicated code

Co-authored-by: Allen Porter <allen.porter@gmail.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Allen Porter <allen.porter@gmail.com>

Author: 3 people

roborock/devices/device_manager.py

@@ -199,7 +199,16 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                 trait = a01.create(product, channel)
             case DeviceVersion.B01:
                 channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
-                trait = b01.create(channel)
+                model_part = product.model.split(".")[-1]
+                if "ss" in model_part:
+                    raise NotImplementedError(
+                        f"Device {device.name} has unsupported version B01_{product.model.strip('.')[-1]}"
+                    )
+                elif "sc" in model_part:
+                    # Q7 devices start with 'sc' in their model naming.
+                    trait = b01.q7.create(channel)
+                else:
+                    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)

roborock/devices/traits/b01/__init__.py

@@ -1,30 +1,5 @@
 """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 roborock.roborock_message import RoborockB01Props
+from .q7 import Q7PropertiesApi
 
-__all__ = [
-    "PropertiesApi",
-]
-
-
-class PropertiesApi(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) -> PropertiesApi:
-    """Create traits for B01 devices."""
-    return PropertiesApi(channel)
+__all__ = ["Q7PropertiesApi", "q7", "q10"]

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

@@ -0,0 +1 @@
+"""Q10"""

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

@@ -0,0 +1,31 @@
+"""Traits for Q7 B01 devices.
+Potentially other devices may fall into this category in the future."""
+
+from roborock.devices.b01_channel import send_decoded_command
+from roborock.devices.mqtt_channel import MqttChannel
+from roborock.devices.traits import Trait
+from roborock.roborock_message import RoborockB01Props
+from roborock.roborock_typing import RoborockB01Q7Methods
+
+__all__ = [
+    "Q7PropertiesApi",
+]
+
+
+class Q7PropertiesApi(Trait):
+    """API for interacting with Q7 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 Q7 properties."""
+        await send_decoded_command(
+            self._channel, dps=10000, command=RoborockB01Q7Methods.GET_PROP, params={"property": props}
+        )
+
+
+def create(channel: MqttChannel) -> Q7PropertiesApi:
+    """Create traits for B01 devices."""
+    return Q7PropertiesApi(channel)

roborock/devices/traits/traits_mixin.py

@@ -31,8 +31,8 @@ class TraitsMixin:
     zeo: a01.ZeoApi | None = None
     """Zeo API, if supported."""
 
-    b01_properties: b01.PropertiesApi | None = None
-    """B01 properties trait, if supported."""
+    b01_q7_properties: b01.Q7PropertiesApi | None = None
+    """B01 Q7 properties trait, if supported."""
 
     def __init__(self, trait: Trait) -> None:
         """Initialize the TraitsMixin with the given trait.

roborock/protocols/b01_protocol.py

@@ -7,7 +7,7 @@
 from Crypto.Cipher import AES
 from Crypto.Util.Padding import pad, unpad
 
-from roborock import RoborockB01Methods
+from roborock import RoborockB01Q7Methods
 from roborock.exceptions import RoborockException
 from roborock.roborock_message import (
     RoborockMessage,
@@ -18,7 +18,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 B01_VERSION = b"B01"
-CommandType = RoborockB01Methods | str
+CommandType = RoborockB01Q7Methods | str
 ParamsType = list | dict | int | None
 
 

roborock/roborock_typing.py

@@ -271,8 +271,8 @@ class RoborockCommand(str, Enum):
     APP_GET_ROBOT_SETTING = "app_get_robot_setting"
 
 
-class RoborockB01Methods(StrEnum):
-    """Methods used by the Roborock B01 model."""
+class RoborockB01Q7Methods(StrEnum):
+    """Methods used by the Roborock Q7 model."""
 
     GET_PROP = "prop.get"
     GET_MAP_LIST = "service.get_map_list"