fix: roborock enum

Author: humbertogontijo

roborock/api.py

@@ -262,7 +262,7 @@ async def get_smart_wash_params(self, device_id: str) -> SmartWashParams:
     async def get_dock_summary(self, device_id: str, dock_type: RoborockDockTypeCode) -> RoborockDockSummary:
         try:
             commands = [self.get_dust_collection_mode(device_id)]
-            if dock_type == RoborockDockTypeCode.EMPTY_WASH_FILL_DOCK:
+            if dock_type == RoborockDockTypeCode['3']:
                 commands += [self.get_wash_towel_mode(device_id), self.get_smart_wash_params(device_id)]
             [
                 dust_collection_mode,
@@ -292,7 +292,7 @@ async def get_prop(self, device_id: str) -> RoborockDeviceProp | None:
                 device_id, clean_summary.records[0]
             )
         dock_summary = None
-        if status and status.dock_type != RoborockDockTypeCode.NO_DOCK:
+        if status and status.dock_type != RoborockDockTypeCode['0']:
             dock_summary = await self.get_dock_summary(device_id, status.dock_type)
         if any([status, dnd_timer, clean_summary, consumable]):
             return RoborockDeviceProp(

roborock/code_mappings.py

@@ -5,10 +5,29 @@
 
 class RoborockEnum(str, Enum):
 
+    def __str__(self):
+        return str(self.value)
+
     @classmethod
     def _missing_(cls, code: int):
         return cls._member_map_.get(str(code))
 
+    @classmethod
+    def as_dict(cls):
+        return {i.name: i.value for i in cls}
+
+    @classmethod
+    def values(cls):
+        return list(cls.as_dict().values())
+
+    @classmethod
+    def keys(cls):
+        return list(cls.as_dict().keys())
+
+    @classmethod
+    def items(cls):
+        return cls.as_dict().items()
+
 
 class RoborockCode:
 

roborock/containers.py

@@ -5,7 +5,8 @@
 
 from dacite import from_dict, Config
 
-from roborock.code_mappings import RoborockDockWashTowelModeCode, RoborockDockTypeCode
+from roborock.code_mappings import RoborockDockWashTowelModeCode, RoborockDockTypeCode, RoborockMopIntensityCode, \
+    RoborockStateCode
 from .code_mappings import RoborockMopModeCode, RoborockDockErrorCode, \
     RoborockErrorCode, RoborockDockDustCollectionModeCode, RoborockFanPowerCode
 
@@ -21,11 +22,7 @@ def decamelize(d):
 
 
 @dataclass
-class RoborockBase(dict):
-    def __init__(self, data: dict[str, any]) -> None:
-        super().__init__()
-        if isinstance(data, dict):
-            self.update(data)
+class RoborockBase:
 
     @classmethod
     def from_dict(cls, data: dict[str, any]):
@@ -164,7 +161,7 @@ class Status(RoborockBase):
     msg_ver: Optional[int] = None
     msg_seq: Optional[int] = None
     status: Optional[str] = None
-    state: Optional[int] = None
+    state: Optional[RoborockStateCode] = None
     battery: Optional[int] = None
     clean_time: Optional[int] = None
     clean_area: Optional[int] = None
@@ -183,7 +180,7 @@ class Status(RoborockBase):
     map_status: Optional[int] = None
     is_locating: Optional[int] = None
     lock_status: Optional[int] = None
-    water_box_mode: Optional[int] = None
+    water_box_mode: Optional[RoborockMopIntensityCode] = None
     mop_intensity: Optional[str] = None
     water_box_carriage_status: Optional[int] = None
     mop_forbidden_enable: Optional[int] = None
@@ -285,12 +282,12 @@ class SmartWashParams(RoborockBase):
 
 @dataclass
 class DustCollectionMode(RoborockBase):
-    mode: RoborockDockDustCollectionModeCode
+    mode: Optional[RoborockDockDustCollectionModeCode] = None
 
 
 @dataclass
 class WashTowelMode(RoborockBase):
-    wash_mode: RoborockDockWashTowelModeCode
+    wash_mode: Optional[RoborockDockWashTowelModeCode] = None
 
 
 @dataclass

tests/test_containers.py

@@ -1,3 +1,6 @@
+from roborock.code_mappings import RoborockStateCode, RoborockErrorCode, RoborockFanPowerCode, RoborockMopIntensityCode, \
+    RoborockDockTypeCode, RoborockMopModeCode, RoborockDockErrorCode
+
 from roborock import UserData, HomeData, Consumable, Status, DNDTimer, CleanSummary, CleanRecord
 from .mock_data import USER_DATA, HOME_DATA_RAW, CONSUMABLE, STATUS, DND_TIMER, CLEAN_SUMMARY, CLEAN_RECORD
 
@@ -95,11 +98,11 @@ def test_status():
     s = Status.from_dict(STATUS)
     assert s.msg_ver == 2
     assert s.msg_seq == 458
-    assert s.state == 8
+    assert s.state == RoborockStateCode['8']
     assert s.battery == 100
     assert s.clean_time == 1176
     assert s.clean_area == 20965000
-    assert s.error_code == 'none'
+    assert s.error_code == RoborockErrorCode['0']
     assert s.map_present == 1
     assert s.in_cleaning == 0
     assert s.in_returning == 0
@@ -109,12 +112,12 @@ def test_status():
     assert s.back_type == -1
     assert s.wash_phase == 0
     assert s.wash_ready == 0
-    assert s.fan_power == 'balanced'
+    assert s.fan_power == RoborockFanPowerCode['102']
     assert s.dnd_enabled == 0
     assert s.map_status == 3
     assert s.is_locating == 0
     assert s.lock_status == 0
-    assert s.water_box_mode == 203
+    assert s.water_box_mode == RoborockMopIntensityCode['203']
     assert s.water_box_carriage_status == 1
     assert s.mop_forbidden_enable == 1
     assert s.camera_status == 3457
@@ -123,15 +126,15 @@ def test_status():
     assert s.home_sec_enable_password == 0
     assert s.adbumper_status == [0, 0, 0]
     assert s.water_shortage_status == 0
-    assert s.dock_type == 'empty_wash_fill_dock'
+    assert s.dock_type == RoborockDockTypeCode['3']
     assert s.dust_collection_status == 0
     assert s.auto_dust_collection == 1
     assert s.avoid_count == 19
-    assert s.mop_mode == 'standard'
+    assert s.mop_mode == RoborockMopModeCode['300']
     assert s.debug_mode == 0
     assert s.collision_avoid_status == 1
     assert s.switch_map_mode == 0
-    assert s.dock_error_status == "ok"
+    assert s.dock_error_status == RoborockDockErrorCode['0']
     assert s.charge_status == 1
     assert s.unsave_map_reason == 0
     assert s.unsave_map_flag == 0