Refactored code to use a single WashInfo command class

Author: lukakama

deebot_client/capabilities.py

@@ -43,15 +43,17 @@
     WorkMode,
     WorkModeEvent,
 )
+from deebot_client.events.wash_info import WashInfoEvent
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
     from _typeshed import DataclassInstance
 
     from deebot_client.command import Command, SetCommand
+    from deebot_client.commands.json.wash_info import SetWashInfo
     from deebot_client.events.efficiency_mode import EfficiencyMode, EfficiencyModeEvent
-    from deebot_client.events.wash_info import WashInfoEvent, WashMode
+    from deebot_client.events.wash_info import WashMode
     from deebot_client.models import CleanAction, CleanMode
 
 
@@ -122,6 +124,14 @@ class CapabilityCleanAction:
     area: Callable[[CleanMode, str, int], Command]
 
 
+@dataclass(frozen=True, kw_only=True)
+class CapabilityWashInfo(CapabilityEvent[WashInfoEvent]):
+    """Capabilities for wash handling."""
+
+    set: Callable[[WashMode | None, int | None], SetWashInfo]
+    wash_modes: tuple[WashMode, ...]
+
+
 @dataclass(frozen=True, kw_only=True)
 class CapabilityClean:
     """Capabilities for clean."""
@@ -132,8 +142,7 @@ class CapabilityClean:
     log: CapabilityEvent[CleanLogEvent] | None = None
     preference: CapabilitySetEnable[CleanPreferenceEvent] | None = None
     work_mode: CapabilitySetTypes[WorkModeEvent, WorkMode] | None = None
-    wash_info_mode: CapabilitySetTypes[WashInfoEvent, WashMode] | None = None
-    wash_info_hot_wash_amount: CapabilitySet[WashInfoEvent, int] | None = None
+    wash_info: CapabilityWashInfo | None = None
 
 
 @dataclass(frozen=True)

deebot_client/commands/json/__init__.py

@@ -40,7 +40,7 @@
 from .true_detect import GetTrueDetect, SetTrueDetect
 from .voice_assistant_state import GetVoiceAssistantState, SetVoiceAssistantState
 from .volume import GetVolume, SetVolume
-from .wash_info import GetWashInfo, SetWashInfoHotWashAmount, SetWashInfoMode
+from .wash_info import GetWashInfo, SetWashInfo
 from .water_info import GetWaterInfo, SetWaterInfo
 from .work_mode import GetWorkMode, SetWorkMode
 
@@ -99,8 +99,7 @@
     "GetVolume",
     "SetVolume",
     "GetWashInfo",
-    "SetWashInfoHotWashAmount",
-    "SetWashInfoMode",
+    "SetWashInfo",
     "GetWaterInfo",
     "SetWaterInfo",
     "GetWorkMode",
@@ -188,8 +187,7 @@
     SetVolume,
 
     GetWashInfo,
-    SetWashInfoHotWashAmount,
-    SetWashInfoMode,
+    SetWashInfo,
 
     GetWaterInfo,
     SetWaterInfo,

deebot_client/commands/json/wash_info.py

@@ -1,23 +1,22 @@
 """WashInfo command module."""
+from __future__ import annotations
 
 from types import MappingProxyType
 from typing import Any
 
 from deebot_client.command import InitParam
-from deebot_client.event_bus import EventBus
-from deebot_client.events import WashInfoEvent, WashMode
-from deebot_client.message import HandlingResult
+from deebot_client.events import WashMode
+from deebot_client.messages.json.wash_info import OnWashInfo
 
 from .common import JsonGetCommand, JsonSetCommand
 
 
-class GetWashInfo(OnWaterInfo, JsonGetCommand):
+class GetWashInfo(OnWashInfo, JsonGetCommand):
     """Get wash info command."""
 
     name = "getWashInfo"
 
 
-
 class SetWashInfo(JsonSetCommand):
     """Set wash info command."""
 
@@ -30,14 +29,17 @@ class SetWashInfo(JsonSetCommand):
         }
     )
 
-    def __init__(self, mode: WashMode | str| None = None, hot_wash_amount: int | None) -> None:
-        args = {}
+    def __init__(
+        self, mode: WashMode | str | None = None, hot_wash_amount: int | None = None
+    ) -> None:
+        args: dict[str, Any] = {}
+
         if isinstance(mode, str):
             mode = WashMode.get(mode)
-        
+
         if mode is not None:
             args["mode"] = mode
-            
+
         if hot_wash_amount is not None:
-            agrs["hot_wash_amount"] = hot_wash_amount
+            args["hot_wash_amount"] = hot_wash_amount
         super().__init__(args)

deebot_client/events/wash_info.py

@@ -1,4 +1,6 @@
 """Wash info event module."""
+from __future__ import annotations
+
 from dataclasses import dataclass
 
 from deebot_client.util import DisplayNameIntEnum
@@ -17,6 +19,6 @@ class WashMode(DisplayNameIntEnum):
 class WashInfoEvent(Event):
     """Wash info event representation."""
 
-    mode: WashMode
+    mode: WashMode | None
     interval: int | None
     hot_wash_amount: int | None

deebot_client/hardware/deebot/p1jij8.py

@@ -15,6 +15,7 @@
     CapabilitySettings,
     CapabilitySetTypes,
     CapabilityStats,
+    CapabilityWashInfo,
 )
 from deebot_client.commands.json.advanced_mode import GetAdvancedMode, SetAdvancedMode
 from deebot_client.commands.json.battery import GetBattery
@@ -53,8 +54,7 @@
 from deebot_client.commands.json.volume import GetVolume, SetVolume
 from deebot_client.commands.json.wash_info import (
     GetWashInfo,
-    SetWashInfoHotWashAmount,
-    SetWashInfoMode,
+    SetWashInfo,
 )
 from deebot_client.commands.json.water_info import GetWaterInfo, SetWaterInfo
 from deebot_client.commands.json.work_mode import GetWorkMode, SetWorkMode
@@ -130,20 +130,15 @@
                     WorkMode.VACUUM_AND_MOP,
                 ),
             ),
-            wash_info_mode=CapabilitySetTypes(
+            wash_info=CapabilityWashInfo(
                 event=WashInfoEvent,
                 get=[GetWashInfo()],
-                set=SetWashInfoMode,
-                types=(
+                set=SetWashInfo,
+                wash_modes=(
                     WashMode.STANDARD,
                     WashMode.HOT,
                 ),
             ),
-            wash_info_hot_wash_amount=CapabilitySet(
-                event=WashInfoEvent,
-                get=[GetWashInfo()],
-                set=SetWashInfoHotWashAmount,
-            ),
         ),
         custom=CapabilityCustomCommand(
             event=CustomCommandEvent, get=[], set=CustomCommand

deebot_client/messages/json/wash_info.py

@@ -1,11 +1,15 @@
 """WashInfo messages."""
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
-from deebot_client.event_bus import EventBus
 from deebot_client.events import WashInfoEvent
 from deebot_client.events.wash_info import WashMode
 from deebot_client.message import HandlingResult, MessageBodyDataDict
 
+if TYPE_CHECKING:
+    from deebot_client.event_bus import EventBus
+
 
 class OnWashInfo(MessageBodyDataDict):
     """On battery message."""
@@ -20,11 +24,15 @@ def _handle_body_data_dict(
 
         :return: A message response
         """
+        mode = data.get("mode")
+        if isinstance(mode, int):
+            mode = WashMode(mode)
+
         event_bus.notify(
             WashInfoEvent(
-                mode=WashMode(int(data["mode"])),
-                hot_wash_amount=data["hot_wash_amount"],
-                interval=data["interval"],
+                mode=mode,
+                hot_wash_amount=data.get("hot_wash_amount"),
+                interval=data.get("interval"),
             )
         )
         return HandlingResult.success()

tests/commands/json/test_wash_info.py

@@ -1,18 +1,21 @@
+from __future__ import annotations
+
 from typing import Any
 
 import pytest
 
 from deebot_client.commands.json import (
     GetWashInfo,
 )
+from deebot_client.commands.json.wash_info import SetWashInfo
 from deebot_client.events import WashInfoEvent, WashMode
 from tests.helpers import (
     get_request_json,
     get_success_body,
     verify_DisplayNameEnum_unique,
 )
 
-from . import assert_command
+from . import assert_command, assert_set_command
 
 
 def test_WashInfo_unique() -> None:
@@ -35,3 +38,22 @@ def test_WashInfo_unique() -> None:
 async def test_GetWashInfo(json: dict[str, Any], expected: WashInfoEvent) -> None:
     json = get_request_json(get_success_body(json))
     await assert_command(GetWashInfo(), json, expected)
+
+
+@pytest.mark.parametrize(("value"), [WashMode.HOT, "hot"])
+async def test_SetWashInfo_mode(value: WashMode | str) -> None:
+    command = SetWashInfo(mode=value)
+    args = {"mode": WashMode.HOT}
+    await assert_set_command(command, args, WashInfoEvent(WashMode.HOT, None, None))
+
+
+def test_SetWashInfo_mode_inexisting_value() -> None:
+    with pytest.raises(ValueError, match="'INEXSTING' is not a valid WashMode member"):
+        SetWashInfo(mode="inexsting")
+
+
+@pytest.mark.parametrize(("value"), [1])
+async def test_SetWashInfo_hot_wash_amount(value: int) -> None:
+    command = SetWashInfo(hot_wash_amount=value)
+    args = {"hot_wash_amount": value}
+    await assert_set_command(command, args, WashInfoEvent(None, None, value))