Fix work mode capability (https://github.com/DeebotUniverse/Deebot-4-…

…Home-Assistant/issues/434)
DeebotUniverse · Oct 25, 2023 · bb01bc2 · bb01bc2
1 parent 34b4662
commit bb01bc2
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/deebot_client/capabilities.py b/deebot_client/capabilities.py
@@ -122,7 +122,7 @@ class CapabilityClean:
     count: CapabilitySet[CleanCountEvent, int] | None = None
     log: CapabilityEvent[CleanLogEvent]
     preference: CapabilitySetEnable[CleanPreferenceEvent] | None = None
-    work_mode: CapabilitySet[WorkModeEvent, WorkMode] | None = None
+    work_mode: CapabilitySetTypes[WorkModeEvent, WorkMode] | None = None
 
 
 @dataclass(frozen=True)

diff --git a/deebot_client/hardware/deebot/p1jij8.py b/deebot_client/hardware/deebot/p1jij8.py
@@ -81,6 +81,7 @@
     VolumeEvent,
     WaterAmount,
     WaterInfoEvent,
+    WorkMode,
     WorkModeEvent,
 )
 from deebot_client.models import StaticDeviceInfo
@@ -106,7 +107,17 @@
             preference=CapabilitySetEnable(
                 CleanPreferenceEvent, [GetCleanPreference()], SetCleanPreference
             ),
-            work_mode=CapabilitySet(WorkModeEvent, [GetWorkMode()], SetWorkMode),
+            work_mode=CapabilitySetTypes(
+                event=WorkModeEvent,
+                get=[GetWorkMode()],
+                set=SetWorkMode,
+                types=(
+                    WorkMode.MOP,
+                    WorkMode.MOP_AFTER_VACUUM,
+                    WorkMode.VACUUM,
+                    WorkMode.VACUUM_AND_MOP,
+                ),
+            ),
         ),
         custom=CapabilityCustomCommand(
             event=CustomCommandEvent, get=[], set=CustomCommand