feat: add more data and region

Author: Lash-L

roborock/data/v1/v1_clean_modes.py

@@ -16,6 +16,8 @@ class VacuumModes(RoborockModeEnum):
     TURBO = ("turbo", 103)
     MAX = ("max", 104)
     MAX_PLUS = ("max_plus", 108)
+    CARPET = ("carpet", 107)
+    OFF_RAISE_MAIN_BRUSH = ("off", 109)
     CUSTOMIZED = ("custom", 106)
     SMART_MODE = ("smart_mode", 110)
 
@@ -45,6 +47,8 @@ class WaterModes(RoborockModeEnum):
     STANDARD = ("standard", 202)
     HIGH = ("high", 203)
     INTENSE = ("intense", 203)
+    MIN = ("min", 205)
+    MAX = ("max", 206)
     CUSTOMIZED = ("custom", 204)
     CUSTOM = ("custom_water_flow", 207)
     EXTREME = ("extreme", 208)
@@ -81,9 +85,14 @@ def get_clean_modes(features: DeviceFeatures) -> list[VacuumModes]:
     if features.is_max_plus_mode_supported or features.is_none_pure_clean_mop_with_max_plus:
         # If the vacuum has max plus mode supported
         modes.append(VacuumModes.MAX_PLUS)
+    if features.is_carpet_deep_clean_supported:
+        modes.append(VacuumModes.CARPET)
     if features.is_pure_clean_mop_supported:
         # If the vacuum is capable of 'pure mop clean' aka no vacuum
-        modes.append(VacuumModes.OFF)
+        if features.is_support_main_brush_up_down_supported:
+            modes.append(VacuumModes.OFF_RAISE_MAIN_BRUSH)
+        else:
+            modes.append(VacuumModes.OFF)
     else:
         # If not, we can add gentle
         modes.append(VacuumModes.GENTLE)

roborock/data/v1/v1_containers.py

@@ -256,6 +256,7 @@ def __repr__(self) -> str:
         return _attr_repr(self)
 
 
+@dataclass
 class StatusV2(RoborockBase):
     """
     This is a new version of the Status object.

roborock/devices/device_manager.py

@@ -215,6 +215,7 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                     web_api,
                     device_cache=device_cache,
                     map_parser_config=map_parser_config,
+                    region=user_data.region,
                 )
             case DeviceVersion.A01:
                 channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)

roborock/devices/traits/v1/__init__.py

@@ -48,7 +48,6 @@
 
 import logging
 from dataclasses import dataclass, field, fields
-from functools import cache
 from typing import Any, get_args
 
 from roborock.data.containers import HomeData, HomeDataProduct, RoborockBase
@@ -174,24 +173,24 @@ def __init__(
         web_api: UserWebApiClient,
         device_cache: DeviceCache,
         map_parser_config: MapParserConfig | None = None,
+        region: str | None = None,
     ) -> None:
         """Initialize the V1TraitProps."""
         self._device_uid = device_uid
         self._rpc_channel = rpc_channel
         self._mqtt_rpc_channel = mqtt_rpc_channel
         self._map_rpc_channel = map_rpc_channel
         self._web_api = web_api
-        self._cache = cache
-        self.device_features = DeviceFeaturesTrait(product.product_nickname, cache)
         self._device_cache = device_cache
+        self._region = region
 
-        self.status = StatusTrait(self.device_features)
+        self.device_features = DeviceFeaturesTrait(product.product_nickname, self._device_cache)
+        self.status = StatusTrait(self.device_features, region=self._region)
         self.consumables = ConsumableTrait()
         self.rooms = RoomsTrait(home_data)
         self.maps = MapsTrait(self.status)
         self.map_content = MapContentTrait(map_parser_config)
         self.home = HomeTrait(self.status, self.maps, self.map_content, self.rooms, self._device_cache)
-        self.device_features = DeviceFeaturesTrait(product.product_nickname, self._device_cache)
         self.network_info = NetworkInfoTrait(device_uid, self._device_cache)
         self.routines = RoutinesTrait(device_uid, web_api)
 
@@ -202,6 +201,8 @@ def __init__(
                 if (union_args := get_args(item.type)) is None or len(union_args) > 0:
                     continue
                 _LOGGER.debug("Trait '%s' is supported, initializing", item.name)
+                if not callable(item.type):
+                    continue
                 trait = item.type()
                 setattr(self, item.name, trait)
             # This is a hack to allow setting the rpc_channel on all traits. This is
@@ -320,6 +321,7 @@ def create(
     web_api: UserWebApiClient,
     device_cache: DeviceCache,
     map_parser_config: MapParserConfig | None = None,
+    region: str | None = None,
 ) -> PropertiesApi:
     """Create traits for V1 devices."""
     return PropertiesApi(
@@ -332,4 +334,5 @@ def create(
         web_api,
         device_cache,
         map_parser_config,
+        region=region,
     )

roborock/devices/traits/v1/status.py

@@ -2,18 +2,22 @@
 from typing import Self
 
 from roborock import CleanRoutes, StatusV2, VacuumModes, WaterModes, get_clean_modes, get_clean_routes, get_water_modes
-from roborock.devices.traits.v1 import DeviceFeaturesTrait, common
 from roborock.roborock_typing import RoborockCommand
 
+from . import common
+from .device_features import DeviceFeaturesTrait
+
 
 class StatusTrait(StatusV2, common.V1TraitMixin):
     """Trait for managing the status of Roborock devices."""
 
     command = RoborockCommand.GET_STATUS
 
-    def __init__(self, device_feature_trait: DeviceFeaturesTrait) -> None:
+    def __init__(self, device_feature_trait: DeviceFeaturesTrait, region: str | None = None) -> None:
         """Initialize the StatusTrait."""
+        super().__init__()
         self._device_features_trait = device_feature_trait
+        self._region = region
 
     @cached_property
     def fan_speed_options(self) -> list[VacuumModes]:
@@ -33,7 +37,7 @@ def water_mode_mapping(self) -> dict[int, str]:
 
     @cached_property
     def mop_route_options(self) -> list[CleanRoutes]:
-        return get_clean_routes(self._device_features_trait, "US")  # TODO: find best place to get region
+        return get_clean_routes(self._device_features_trait, self._region or "us")
 
     @cached_property
     def mop_route_mapping(self) -> dict[int, str]:

tests/devices/__snapshots__/test_v1_device.ambr

@@ -832,7 +832,7 @@
   })
 # ---
 # name: test_device_trait_command_parsing[status]
-  StatusTrait(adbumper_status=None, auto_dust_collection=None, avoid_count=None, back_type=None, battery=100, camera_status=None, charge_status=None, clean_area=91287500, clean_fluid_status=None, clean_percent=None, clean_time=5405, clear_water_box_status=None, collision_avoid_status=None, command=<RoborockCommand.GET_STATUS: 'get_status'>, common_status=None, corner_clean_mode=None, current_map=0, debug_mode=None, dirty_water_box_status=None, distance_off=0, dnd_enabled=1, dock_cool_fan_status=None, dock_error_status=None, dock_type=None, dry_status=None, dss=None, dust_bag_status=None, dust_collection_status=None, error_code=<RoborockErrorCode.none: 0>, error_code_name='none', fan_power=<RoborockFanSpeedS7MaxV.custom: 106>, fan_power_name='custom', fan_power_options=['off', 'quiet', 'balanced', 'turbo', 'max', 'custom', 'max_plus'], hatch_door_status=None, home_sec_enable_password=None, home_sec_status=None, in_cleaning=<RoborockInCleaning.complete: 0>, in_fresh_state=1, in_returning=0, in_warmup=None, is_exploring=None, is_locating=0, kct=None, lab_status=1, last_clean_t=None, lock_status=0, map_present=1, map_status=3, mop_forbidden_enable=0, mop_mode=None, mop_mode_name=None, msg_seq=515, msg_ver=2, rdt=None, repeat=None, replenish_mode=None, rss=None, square_meter_clean_area=91.3, state=<RoborockStateCode.charging: 8>, state_name='charging', subdivision_sets=None, switch_map_mode=None, unsave_map_flag=0, unsave_map_reason=4, wash_phase=None, wash_ready=None, wash_status=None, water_box_carriage_status=0, water_box_filter_status=None, water_box_mode=<RoborockMopIntensityS7.custom: 204>, water_box_mode_name='custom', water_box_status=0, water_shortage_status=None)
+  StatusTrait(adbumper_status=None, auto_dust_collection=None, avoid_count=None, back_type=None, battery=100, camera_status=None, charge_status=None, clean_area=91287500, clean_fluid_status=None, clean_percent=None, clean_time=5405, clear_water_box_status=None, collision_avoid_status=None, command=<RoborockCommand.GET_STATUS: 'get_status'>, common_status=None, corner_clean_mode=None, current_map=0, debug_mode=None, dirty_water_box_status=None, distance_off=0, dnd_enabled=1, dock_cool_fan_status=None, dock_error_status=None, dock_type=None, dry_status=None, dss=None, dust_bag_status=None, dust_collection_status=None, error_code=<RoborockErrorCode.none: 0>, error_code_name='none', fan_power=106, fan_speed_mapping={101: 'QUIET', 102: 'BALANCED', 103: 'TURBO', 104: 'MAX', 105: 'GENTLE'}, fan_speed_name=None, fan_speed_options=[<VacuumModes.QUIET: 'quiet'>, <VacuumModes.BALANCED: 'balanced'>, <VacuumModes.TURBO: 'turbo'>, <VacuumModes.MAX: 'max'>, <VacuumModes.GENTLE: 'gentle'>], hatch_door_status=None, home_sec_enable_password=None, home_sec_status=None, in_cleaning=<RoborockInCleaning.complete: 0>, in_fresh_state=1, in_returning=0, in_warmup=None, is_exploring=None, is_locating=0, kct=None, lab_status=1, last_clean_t=None, lock_status=0, map_present=1, map_status=3, mop_forbidden_enable=0, mop_mode=None, mop_route_mapping={300: 'STANDARD', 301: 'DEEP'}, mop_route_name=None, mop_route_options=[<CleanRoutes.STANDARD: 'standard'>, <CleanRoutes.DEEP: 'deep'>], msg_seq=515, msg_ver=2, rdt=None, repeat=None, replenish_mode=None, rss=None, square_meter_clean_area=91.3, state=<RoborockStateCode.charging: 8>, state_name='charging', subdivision_sets=None, switch_map_mode=None, unsave_map_flag=0, unsave_map_reason=4, wash_phase=None, wash_ready=None, wash_status=None, water_box_carriage_status=0, water_box_filter_status=None, water_box_mode=204, water_box_status=0, water_mode_mapping={200: 'OFF', 201: 'LOW', 202: 'MEDIUM', 203: 'HIGH'}, water_mode_name=None, water_mode_options=[<WaterModes.OFF: 'off'>, <WaterModes.LOW: 'low'>, <WaterModes.MEDIUM: 'medium'>, <WaterModes.HIGH: 'high'>], water_shortage_status=None)
 # ---
 # name: test_device_trait_command_parsing[status].1
   dict({

tests/devices/test_v1_device.py

@@ -65,6 +65,7 @@ def device_fixture(channel: AsyncMock, rpc_channel: AsyncMock, mqtt_rpc_channel:
             AsyncMock(),
             AsyncMock(),
             device_cache=DeviceCache(HOME_DATA.devices[0].duid, NoCache()),
+            region=USER_DATA.region,
         ),
     )
 

tests/devices/traits/v1/fixtures.py

@@ -81,6 +81,7 @@ def device_fixture(
             mock_map_rpc_channel,
             web_api_client,
             device_cache=device_cache,
+            region=USER_DATA.region,
         ),
     )