Add StatusField support for DeviceFeatures

.pre-commit-config.yaml

@@ -42,7 +42,7 @@ repos:
         hooks:
             -   id: mypy
                 exclude: cli.py
-                additional_dependencies: [ "types-paho-mqtt" ]
+                additional_dependencies: [ "types-paho-mqtt", "types-pyyaml" ]
     -   repo: https://github.com/alessandrojcm/commitlint-pre-commit-hook
         rev: v9.23.0
         hooks:

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,20 +115,25 @@ 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):
     """This status will be deprecated in favor of StatusV2."""
 
     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
@@ -139,12 +146,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
@@ -162,15 +169,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/device_features.py

@@ -1,8 +1,6 @@
-from __future__ import annotations
-
 from dataclasses import dataclass, field, fields
 from enum import IntEnum, StrEnum
-from typing import Any
+from typing import Any, Self
 
 from roborock.data.code_mappings import RoborockProductNickname
 from roborock.data.containers import RoborockBase
@@ -566,7 +564,7 @@ def from_feature_flags(
         new_feature_info_str: str,
         feature_info: list[int],
         product_nickname: RoborockProductNickname | None,
-    ) -> DeviceFeatures:
+    ) -> Self:
         """Creates a DeviceFeatures instance from raw feature flags.
         :param new_feature_info: A int from get_init_status (sometimes can be found in homedata, but it is not always)
         :param new_feature_info_str: A hex string from get_init_status or home_data.

roborock/devices/traits/v1/device_features.py

@@ -40,10 +40,11 @@ def __init__(self, product: HomeDataProduct, device_cache: DeviceCache) -> None:
         self.converter = DeviceTraitsConverter(product)
         self._product = product
         self._device_cache = device_cache
-        # All fields of DeviceFeatures are required. Initialize them to False
+        # All boolean fields of DeviceFeatures are required. Initialize them to False
         # so we have some known state.
         for field in fields(self):
-            setattr(self, field.name, False)
+            if field.type is bool:
+                setattr(self, field.name, False)
 
     def is_field_supported(self, cls: type[RoborockBase], field_name: FieldNameBase) -> bool:
         """Determines if the specified field is supported by this device.
@@ -61,11 +62,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': True,
     'dry_status': True,
     'fan_power': True,
     'state': True,
@@ -13,6 +14,7 @@
   dict({
     'battery': True,
     'charge_status': True,
+    'clean_percent': True,
     'dry_status': True,
     'fan_power': True,
     'state': True,
@@ -23,6 +25,7 @@
   dict({
     'battery': True,
     'charge_status': True,
+    'clean_percent': True,
     'dry_status': True,
     'fan_power': True,
     'state': True,

tests/devices/traits/v1/fixtures.py

@@ -1,6 +1,7 @@
 """Fixtures for V1 trait tests."""
 
 from copy import deepcopy
+from typing import Any
 from unittest.mock import AsyncMock
 
 import pytest
@@ -107,16 +108,34 @@ 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 model-specific APP_GET_INIT_STATUS data.
+
+    Uses real device feature data from device_info.yaml when available for the
+    product model, falling back to the default mock data otherwise.
+    """
+    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")
+    device_info_data = mock_data.DEVICE_INFO.get(product.model, {})
+    return {
+        **mock_data.APP_GET_INIT_STATUS,
+        **device_info_data,
+    }
+
+
 @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 keyed by product model.
+# Each entry contains the fields needed for APP_GET_INIT_STATUS responses.
+_DEVICE_INFO_DATA = yaml.safe_load(pathlib.Path("device_info.yaml").read_text())
+DEVICE_INFO: dict[str, dict[str, Any]] = {
+    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,