Add StatusField support for DeviceFeatures

Author: allenporter

device_info.yaml

@@ -1,61 +1,68 @@
 roborock.vacuum.a15:
-  Protocol Version: '1.0'
+  Protocol Version: "1.0"
   Product Nickname: TANOSS
   New Feature Info: 636084721975295
-  New Feature Info Str: '0000000000002000'
+  New Feature Info Str: "0000000000002000"
   Feature Info:
-  - 111
-  - 112
-  - 113
-  - 114
-  - 115
-  - 116
-  - 117
-  - 118
-  - 119
-  - 120
-  - 122
-  - 123
-  - 124
-  - 125
+    - 111
+    - 112
+    - 113
+    - 114
+    - 115
+    - 116
+    - 117
+    - 118
+    - 119
+    - 120
+    - 122
+    - 123
+    - 124
+    - 125
 roborock.vacuum.a87:
-  Protocol Version: '1.0'
+  Protocol Version: "1.0"
   Product Nickname: PEARLPLUS
   New Feature Info: 4499197267967999
   New Feature Info Str: 508A977F7EFEFFFF
   Feature Info:
-  - 111
-  - 112
-  - 113
-  - 114
-  - 115
-  - 116
-  - 117
-  - 118
-  - 119
-  - 120
-  - 121
-  - 122
-  - 123
-  - 124
-  - 125
+    - 111
+    - 112
+    - 113
+    - 114
+    - 115
+    - 116
+    - 117
+    - 118
+    - 119
+    - 120
+    - 121
+    - 122
+    - 123
+    - 124
+    - 125
 roborock.vacuum.s5e:
-  Protocol Version: '1.0'
+  Protocol Version: "1.0"
   Product Nickname: RUBYSLITE
   New Feature Info: 633887780925447
-  New Feature Info Str: '0000000000002000'
+  New Feature Info Str: "0000000000002000"
   Feature Info:
-  - 111
-  - 112
-  - 113
-  - 114
-  - 115
-  - 116
-  - 117
-  - 118
-  - 119
-  - 120
-  - 122
-  - 123
-  - 124
-  - 125
+    - 111
+    - 112
+    - 113
+    - 114
+    - 115
+    - 116
+    - 117
+    - 118
+    - 119
+    - 120
+    - 122
+    - 123
+    - 124
+    - 125
+roborock.vacuum.a144:
+  Protocol Version: "1.0"
+  Product Nickname: CORALPRO
+  New Feature Info: 4499197267967999
+  New Feature Info Str: "0000000000ED3EDDCFFF8F7F7EFEFFFF"
+  # We don't have this value populated yet
+  Feature Info: []

roborock/data/v1/v1_containers.py

@@ -103,8 +103,10 @@ class StatusField(FieldNameBase):
     to understand if a feature is supported by the device using `is_field_supported`.
 
     The enum values are names of fields in the `Status` class. Each field is
-    annotated with `requires_schema_code` metadata to map the field to a schema
-    code in the product schema, which may have a different name than the field/attribute name.
+    annotated with one of the following:
+    - `requires_schema_code` metadata to map the field to a schema code in the
+       product schema, which may have a different name than the field/attribute name.
+    - `requires_supported_feature` metadata to map the field to a field in `DeviceFeatures`.
     """
 
     STATE = "state"
@@ -113,18 +115,23 @@ class StatusField(FieldNameBase):
     WATER_BOX_MODE = "water_box_mode"
     CHARGE_STATUS = "charge_status"
     DRY_STATUS = "dry_status"
+    CLEAN_PERCENT = "clean_percent"
 
 
-def _requires_schema_code(requires_schema_code: str, default=None) -> Any:
+def _requires_schema_code(requires_schema_code: str, default: Any = None) -> Any:
     return field(metadata={"requires_schema_code": requires_schema_code}, default=default)
 
 
+def _requires_supported_feature(requires_supported_feature: str, default: Any = None) -> Any:
+    return field(metadata={"requires_supported_feature": requires_supported_feature}, default=default)
+
+
 @dataclass
 class Status(RoborockBase):
     msg_ver: int | None = None
     msg_seq: int | None = None
-    state: RoborockStateCode | None = _requires_schema_code("state", default=None)
-    battery: int | None = _requires_schema_code("battery", default=None)
+    state: RoborockStateCode | None = _requires_schema_code("state")
+    battery: int | None = _requires_schema_code("battery")
     clean_time: int | None = None
     clean_area: int | None = None
     error_code: RoborockErrorCode | None = None
@@ -137,12 +144,12 @@ class Status(RoborockBase):
     back_type: int | None = None
     wash_phase: int | None = None
     wash_ready: int | None = None
-    fan_power: RoborockFanPowerCode | None = _requires_schema_code("fan_power", default=None)
+    fan_power: RoborockFanPowerCode | None = _requires_schema_code("fan_power")
     dnd_enabled: int | None = None
     map_status: int | None = None
     is_locating: int | None = None
     lock_status: int | None = None
-    water_box_mode: RoborockMopIntensityCode | None = _requires_schema_code("water_box_mode", default=None)
+    water_box_mode: RoborockMopIntensityCode | None = _requires_schema_code("water_box_mode")
     water_box_carriage_status: int | None = None
     mop_forbidden_enable: int | None = None
     camera_status: int | None = None
@@ -160,15 +167,15 @@ class Status(RoborockBase):
     collision_avoid_status: int | None = None
     switch_map_mode: int | None = None
     dock_error_status: RoborockDockErrorCode | None = None
-    charge_status: int | None = _requires_schema_code("charge_status", default=None)
+    charge_status: int | None = _requires_schema_code("charge_status")
     unsave_map_reason: int | None = None
     unsave_map_flag: int | None = None
     wash_status: int | None = None
     distance_off: int | None = None
     in_warmup: int | None = None
-    dry_status: int | None = _requires_schema_code("drying_status", default=None)
+    dry_status: int | None = _requires_schema_code("drying_status")
     rdt: int | None = None
-    clean_percent: int | None = None
+    clean_percent: int | None = _requires_supported_feature("is_support_clean_estimate")
     rss: int | None = None
     dss: int | None = None
     common_status: int | None = None

roborock/devices/traits/v1/device_features.py

@@ -39,11 +39,12 @@ def is_field_supported(self, cls: type[RoborockBase], field_name: FieldNameBase)
             raise ValueError(f"Field {field_name} not found in {cls}")
 
         requires_schema_code = dataclass_field.metadata.get("requires_schema_code", None)
-        if requires_schema_code is None:
-            # We assume the field is supported
-            return True
-        # If the field requires a protocol that is not supported, we return False
-        return requires_schema_code in self._product.supported_schema_codes
+        if requires_schema_code is not None:
+            return requires_schema_code in self._product.supported_schema_codes
+        requires_supported_feature = dataclass_field.metadata.get("requires_supported_feature", None)
+        if requires_supported_feature is not None:
+            return getattr(self, requires_supported_feature)
+        return True
 
     async def refresh(self) -> None:
         """Refresh the contents of this trait.

tests/devices/traits/v1/__snapshots__/test_device_features.ambr

@@ -3,6 +3,7 @@
   dict({
     'battery': True,
     'charge_status': True,
+    'clean_percent': False,
     'dry_status': True,
     'fan_power': True,
     'state': True,
@@ -13,6 +14,7 @@
   dict({
     'battery': True,
     'charge_status': True,
+    'clean_percent': False,
     'dry_status': True,
     'fan_power': True,
     'state': True,
@@ -23,6 +25,7 @@
   dict({
     'battery': True,
     'charge_status': True,
+    'clean_percent': False,
     'dry_status': True,
     'fan_power': True,
     'state': True,

tests/devices/traits/v1/fixtures.py

@@ -1,5 +1,6 @@
 """Fixtures for V1 trait tests."""
 
+from typing import Any
 from unittest.mock import AsyncMock
 
 import pytest
@@ -105,16 +106,26 @@ def dock_type_code_fixture(request: pytest.FixtureRequest) -> RoborockDockTypeCo
     return RoborockDockTypeCode.s7_max_ultra_dock
 
 
+@pytest.fixture(name="mock_app_get_init_status")
+def mock_app_get_init_status_fixture(device_info: HomeDataDevice, products: list[HomeDataProduct]) -> dict[str, Any]:
+    """Fixture to provide a DeviceFeaturesInfo instance for tests."""
+    product = next(filter(lambda product: product.id == device_info.product_id, products))
+    if not product:
+        raise ValueError(f"Product {device_info.product_id} not found")
+    return mock_data.APP_GET_INIT_STATUS
+
+
 @pytest.fixture(autouse=True)
 async def discover_features_fixture(
     device: RoborockDevice,
     mock_rpc_channel: AsyncMock,
+    mock_app_get_init_status: dict[str, Any],
     dock_type_code: RoborockDockTypeCode | None,
 ) -> None:
     """Fixture to handle device feature discovery."""
     assert device.v1_properties
     mock_rpc_channel.send_command.side_effect = [
-        [mock_data.APP_GET_INIT_STATUS],
+        [mock_app_get_init_status],
         {
             **mock_data.STATUS,
             "dock_type": dock_type_code,

tests/mock_data.py

@@ -5,6 +5,8 @@
 import pathlib
 from typing import Any
 
+import yaml
+
 # All data is based on a U.S. customer with a Roborock S7 MaxV Ultra
 USER_EMAIL = "user@domain.com"
 
@@ -141,6 +143,18 @@
 ZEO_ONE_DEVICE_DATA = DEVICES["home_data_device_zeo_one.json"]
 SAROS_10R_DEVICE_DATA = DEVICES["home_data_device_saros_10r.json"]
 
+# Additional Device Features info from YAML. In the future we can merge these
+# all into a similar format or get from diagnostrics.
+_DEVICE_INFO_DATA = yaml.safe_load(pathlib.Path("device_info.yaml").read_text())
+DEVICE_INFO = {
+    product_model: {
+        "new_feature_info": data.get("New Feature Info"),
+        "new_feature_info_str": data.get("New Feature Info Str"),
+        "feature_info": data.get("Feature Info"),
+    }
+    for product_model, data in _DEVICE_INFO_DATA.items()
+}
+
 
 HOME_DATA_RAW: dict[str, Any] = {
     "id": 123456,