Add sweep_type to WaterInfo event (#511)

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: Robert Resch <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

deebot_client/command.py

@@ -280,6 +280,7 @@ class InitParam:
 
     type_: type
     name: str | None = None
+    optional: bool = field(default=False, kw_only=True)
 
 
 class CommandMqttP2P(Command, ABC):
@@ -303,7 +304,12 @@ def create_from_mqtt(cls, data: dict[str, Any]) -> CommandMqttP2P:
                 # Remove field
                 data.pop(name, None)
             else:
-                values[param.name or name] = _pop_or_raise(name, param.type_, data)
+                try:
+                    values[param.name or name] = _pop_or_raise(name, param.type_, data)
+                except KeyError as err:
+                    if not param.optional:
+                        msg = f'"{name}" is missing in {data}'
+                        raise DeebotError(msg) from err
 
         if data:
             _LOGGER.debug("Following data will be ignored: %s", data)
@@ -312,11 +318,7 @@ def create_from_mqtt(cls, data: dict[str, Any]) -> CommandMqttP2P:
 
 
 def _pop_or_raise(name: str, type_: type, data: dict[str, Any]) -> Any:
-    try:
-        value = data.pop(name)
-    except KeyError as err:
-        msg = f'"{name}" is missing in {data}'
-        raise DeebotError(msg) from err
+    value = data.pop(name)
     try:
         return type_(value)
     except ValueError as err:

deebot_client/commands/json/water_info.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING, Any
 
 from deebot_client.command import InitParam
-from deebot_client.events import WaterAmount, WaterInfoEvent
+from deebot_client.events import SweepType, WaterAmount, WaterInfoEvent
 from deebot_client.message import HandlingResult
 from deebot_client.util import get_enum
 
@@ -33,8 +33,15 @@ def _handle_body_data_dict(
         if mop_attached is not None:
             mop_attached = bool(mop_attached)
 
+        if sweep_type := data.get("sweepType"):
+            sweep_type = SweepType(int(sweep_type))
+
         event_bus.notify(
-            WaterInfoEvent(WaterAmount(int(data["amount"])), mop_attached=mop_attached)
+            WaterInfoEvent(
+                WaterAmount(int(data["amount"])),
+                sweep_type,
+                mop_attached=mop_attached,
+            )
         )
         return HandlingResult.success()
 
@@ -48,10 +55,19 @@ class SetWaterInfo(JsonSetCommand):
         {
             "amount": InitParam(WaterAmount),
             "enable": None,  # Remove it as we don't can set it (App includes it)
+            "sweepType": InitParam(SweepType, "sweep_type", optional=True),
         }
     )
 
-    def __init__(self, amount: WaterAmount | str) -> None:
+    def __init__(
+        self, amount: WaterAmount | str, sweep_type: SweepType | str | None = None
+    ) -> None:
+        params = {}
         if isinstance(amount, str):
             amount = get_enum(WaterAmount, amount)
-        super().__init__({"amount": amount.value})
+        params["amount"] = amount.value
+        if sweep_type:
+            if isinstance(sweep_type, str):
+                sweep_type = get_enum(SweepType, sweep_type)
+            params["sweepType"] = sweep_type.value
+        super().__init__(params)

deebot_client/events/__init__.py

@@ -24,7 +24,7 @@
     PositionType,
 )
 from .network import NetworkInfoEvent
-from .water_info import WaterAmount, WaterInfoEvent
+from .water_info import SweepType, WaterAmount, WaterInfoEvent
 from .work_mode import WorkMode, WorkModeEvent
 
 if TYPE_CHECKING:
@@ -52,6 +52,7 @@
     "PositionType",
     "PositionsEvent",
     "SweepModeEvent",
+    "SweepType",
     "WaterAmount",
     "WaterInfoEvent",
     "WorkMode",

deebot_client/events/water_info.py

@@ -18,10 +18,19 @@ class WaterAmount(IntEnum):
     ULTRAHIGH = 4
 
 
+@unique
+class SweepType(IntEnum):
+    """Enum class for all possible sweeping types."""
+
+    STANDARD = 1
+    DEEP = 2
+
+
 @dataclass(frozen=True)
 class WaterInfoEvent(Event):
     """Water info event representation."""
 
     amount: WaterAmount
     # None means no data available
+    sweep_type: SweepType | None = None
     mop_attached: bool | None = field(kw_only=True, default=None)

tests/commands/json/test_water_info.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
+import re
 from typing import Any
 
 import pytest
 
 from deebot_client.commands.json import GetWaterInfo, SetWaterInfo
-from deebot_client.events import WaterAmount, WaterInfoEvent
+from deebot_client.events import SweepType, WaterAmount, WaterInfoEvent
 from tests.helpers import (
     get_request_json,
     get_success_body,
@@ -26,22 +27,63 @@
             {"amount": 4, "enable": 0},
             WaterInfoEvent(WaterAmount.ULTRAHIGH, mop_attached=False),
         ),
+        (
+            {"amount": 4, "sweepType": 1, "enable": 0},
+            WaterInfoEvent(
+                WaterAmount.ULTRAHIGH, SweepType.STANDARD, mop_attached=False
+            ),
+        ),
+        (
+            {"amount": 4, "sweepType": 2, "enable": 0},
+            WaterInfoEvent(WaterAmount.ULTRAHIGH, SweepType.DEEP, mop_attached=False),
+        ),
     ],
 )
 async def test_GetWaterInfo(json: dict[str, Any], expected: WaterInfoEvent) -> None:
     json = get_request_json(get_success_body(json))
     await assert_command(GetWaterInfo(), json, expected)
 
 
-@pytest.mark.parametrize(("value"), [WaterAmount.MEDIUM, "medium"])
-async def test_SetWaterInfo(value: WaterAmount | str) -> None:
-    command = SetWaterInfo(value)
+@pytest.mark.parametrize(("water_value"), [WaterAmount.MEDIUM, "medium"])
+@pytest.mark.parametrize(("sweep_value"), [SweepType.STANDARD, "standard", None])
+async def test_SetWaterInfo_Wateramount(
+    water_value: WaterAmount | str, sweep_value: SweepType | str | None
+) -> None:
+    command = SetWaterInfo(water_value, sweep_value)
     args = {"amount": 2}
-    await assert_set_command(command, args, WaterInfoEvent(WaterAmount.MEDIUM))
+    if sweep_value:
+        args["sweepType"] = 1
+    await assert_set_command(
+        command,
+        args,
+        WaterInfoEvent(WaterAmount.MEDIUM, SweepType.STANDARD if sweep_value else None),
+    )
 
 
-def test_SetWaterInfo_inexisting_value() -> None:
-    with pytest.raises(
-        ValueError, match="'INEXSTING' is not a valid WaterAmount member"
-    ):
-        SetWaterInfo("inexsting")
+@pytest.mark.parametrize(
+    ("command_values", "error", "error_message"),
+    [
+        (
+            {"bla": "inexsting"},
+            TypeError,
+            re.escape(
+                "SetWaterInfo.__init__() got an unexpected keyword argument 'bla'"
+            ),
+        ),
+        (
+            {"amount": "inexsting"},
+            ValueError,
+            "'INEXSTING' is not a valid WaterAmount member",
+        ),
+        (
+            {"amount": WaterAmount.HIGH, "sweep_type": "inexsting"},
+            ValueError,
+            "'INEXSTING' is not a valid SweepType member",
+        ),
+    ],
+)
+def test_SetWaterInfo_inexisting_value(
+    command_values: dict[str, Any], error: type[Exception], error_message: str
+) -> None:
+    with pytest.raises(error, match=error_message):
+        SetWaterInfo(**command_values)