Add WaterCustomAmountEvent and updated water capability

deebot_client/capabilities.py

@@ -125,6 +125,14 @@ class CapabilitySetTypes[E: Event, **P, T](CapabilitySet[E, P], CapabilityTypes[
     """Capability for set command and types."""
 
 
+@dataclass(frozen=True, kw_only=True)
+class CapabilityNumber[E: Event, **P](CapabilitySet[E, P]):
+    """Capability for a number entity with min and max."""
+
+    min: int
+    max: int
+
+
 @dataclass(frozen=True, kw_only=True)
 class CapabilityCleanAction:
     """Capabilities for clean action."""
@@ -228,11 +236,14 @@ class CapabilityStation:
 class CapabilityWater:
     """Capabilities for water."""
 
-    amount: CapabilitySetTypes[
-        water_info.WaterAmountEvent,
-        [water_info.WaterAmount | str],
-        water_info.WaterAmount,
-    ]
+    amount: (
+        CapabilitySetTypes[
+            water_info.WaterAmountEvent,
+            [water_info.WaterAmount | str],
+            water_info.WaterAmount,
+        ]
+        | CapabilityNumber[water_info.WaterCustomAmountEvent, [int]]
+    )
     mop_attached: CapabilityEvent[water_info.MopAttachedEvent]
 
 

deebot_client/commands/json/water_info.py

@@ -11,6 +11,7 @@
     SweepType,
     WaterAmount,
     WaterAmountEvent,
+    WaterCustomAmountEvent,
     WaterSweepTypeEvent,
 )
 from deebot_client.message import HandlingResult
@@ -35,7 +36,11 @@ def _handle_body_data_dict(
 
         :return: A message response
         """
-        event_bus.notify(WaterAmountEvent(WaterAmount(int(data["amount"]))))
+        if "amount" in data:
+            event_bus.notify(WaterAmountEvent(WaterAmount(int(data["amount"]))))
+
+        if "customAmount" in data:
+            event_bus.notify(WaterCustomAmountEvent(int(data["customAmount"])))
 
         if (mop_attached := data.get("enable")) is not None:
             event_bus.notify(MopAttachedEvent(bool(mop_attached)))
@@ -53,19 +58,32 @@ class SetWaterInfo(JsonSetCommand):
     get_command = GetWaterInfo
     _mqtt_params = MappingProxyType(
         {
-            "amount": InitParam(WaterAmount),
+            "amount": InitParam(WaterAmount, optional=True),
+            "customAmount": InitParam(int, "custom_amount", optional=True),
             "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, sweep_type: SweepType | str | None = None
+        self,
+        amount: WaterAmount | str | None = None,
+        custom_amount: int | None = None,
+        sweep_type: SweepType | str | None = None,
     ) -> None:
         params = {}
-        if isinstance(amount, str):
-            amount = get_enum(WaterAmount, amount)
-        params["amount"] = amount.value
+        if amount is not None:
+            if custom_amount is not None:
+                raise ValueError("Only one of amount or custom_amount can be provided.")
+
+            if isinstance(amount, str):
+                amount = get_enum(WaterAmount, amount)
+            params["amount"] = amount.value
+        elif custom_amount is not None:
+            params["customAmount"] = custom_amount
+        else:
+            raise ValueError("Either amount or custom_amount must be provided.")
+
         if sweep_type:
             if isinstance(sweep_type, str):
                 sweep_type = get_enum(SweepType, sweep_type)

deebot_client/events/water_info.py

@@ -11,6 +11,7 @@
     "SweepType",
     "WaterAmount",
     "WaterAmountEvent",
+    "WaterCustomAmountEvent",
     "WaterSweepTypeEvent",
 ]
 
@@ -37,6 +38,10 @@ class WaterAmountEvent(ValueEvent[WaterAmount]):
     """Water amount event."""
 
 
+class WaterCustomAmountEvent(ValueEvent[int]):
+    """Water custom amount event."""
+
+
 class WaterSweepTypeEvent(ValueEvent[SweepType]):
     """Water sweep type event."""
 

deebot_client/hardware/deebot/xco2fc.py

@@ -11,6 +11,7 @@
     CapabilityExecute,
     CapabilityLifeSpan,
     CapabilityMap,
+    CapabilityNumber,
     CapabilitySet,
     CapabilitySetEnable,
     CapabilitySettings,
@@ -190,16 +191,12 @@
             total=CapabilityEvent(TotalStatsEvent, [GetTotalStats()]),
         ),
         water=CapabilityWater(
-            amount=CapabilitySetTypes(
-                event=water_info.WaterAmountEvent,
+            amount=CapabilityNumber(
+                event=water_info.WaterCustomAmountEvent,
                 get=[GetWaterInfo()],
-                set=SetWaterInfo,
-                types=(
-                    water_info.WaterAmount.LOW,
-                    water_info.WaterAmount.MEDIUM,
-                    water_info.WaterAmount.HIGH,
-                    water_info.WaterAmount.ULTRAHIGH,
-                ),
+                set=lambda custom_amount: SetWaterInfo(custom_amount=custom_amount),
+                min=0,
+                max=50,
             ),
             mop_attached=CapabilityEvent(water_info.MopAttachedEvent, [GetWaterInfo()]),
         ),

tests/commands/json/test_water_info.py

@@ -11,6 +11,7 @@
     SweepType,
     WaterAmount,
     WaterAmountEvent,
+    WaterCustomAmountEvent,
     WaterSweepTypeEvent,
 )
 from tests.helpers import (
@@ -58,26 +59,74 @@
                 WaterSweepTypeEvent(SweepType.DEEP),
             ),
         ),
+        (
+            {
+                "customAmount": 30,
+                "enable": 1,
+                "mopCount": 2,
+                "sideMop": 0,
+                "sweepType": 1,
+                "type": 1,
+            },
+            (
+                WaterCustomAmountEvent(30),
+                MopAttachedEvent(True),
+                WaterSweepTypeEvent(SweepType.STANDARD),
+            ),
+        ),
     ],
 )
 async def test_GetWaterInfo(json: dict[str, Any], expected: tuple[Event, ...]) -> None:
     json, firmware_event = get_request_json(get_success_body(json))
     await assert_command(GetWaterInfo(), json, (firmware_event, *expected))
 
 
-@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
+@pytest.mark.parametrize(
+    ("command", "args", "expected_events"),
+    [
+        (
+            SetWaterInfo(WaterAmount.MEDIUM),
+            {"amount": 2},
+            [WaterAmountEvent(WaterAmount.MEDIUM)],
+        ),
+        (SetWaterInfo("high"), {"amount": 3}, [WaterAmountEvent(WaterAmount.HIGH)]),
+        (
+            SetWaterInfo(WaterAmount.LOW, sweep_type=SweepType.STANDARD),
+            {"amount": 1, "sweepType": 1},
+            [
+                WaterAmountEvent(WaterAmount.LOW),
+                WaterSweepTypeEvent(SweepType.STANDARD),
+            ],
+        ),
+        (
+            SetWaterInfo(WaterAmount.ULTRAHIGH, sweep_type="deep"),
+            {"amount": 4, "sweepType": 2},
+            [
+                WaterAmountEvent(WaterAmount.ULTRAHIGH),
+                WaterSweepTypeEvent(SweepType.DEEP),
+            ],
+        ),
+        (
+            SetWaterInfo(custom_amount=30),
+            {"customAmount": 30},
+            [WaterCustomAmountEvent(30)],
+        ),
+        (
+            SetWaterInfo(custom_amount=30, sweep_type="deep"),
+            {"customAmount": 30, "sweepType": 2},
+            [
+                WaterCustomAmountEvent(30),
+                WaterSweepTypeEvent(SweepType.DEEP),
+            ],
+        ),
+    ],
+)
+async def test_SetWaterInfo(
+    command: SetWaterInfo,
+    args: dict[str, Any],
+    expected_events: list[Event],
 ) -> None:
-    command = SetWaterInfo(water_value, sweep_value)
-    args = {"amount": 2}
-    expected_events: list[Event] = [
-        WaterAmountEvent(WaterAmount.MEDIUM),
-    ]
-    if sweep_value:
-        args["sweepType"] = 1
-        expected_events.append(WaterSweepTypeEvent(SweepType.STANDARD))
+    """Test SetWaterInfo."""
     await assert_set_command(command, args, expected_events)
 
 
@@ -101,9 +150,19 @@ async def test_SetWaterInfo_Wateramount(
             ValueError,
             "'INEXSTING' is not a valid SweepType member",
         ),
+        (
+            {},
+            ValueError,
+            "Either amount or custom_amount must be provided.",
+        ),
+        (
+            {"amount": WaterAmount.ULTRAHIGH, "custom_amount": "40"},
+            ValueError,
+            "Only one of amount or custom_amount can be provided.",
+        ),
     ],
 )
-def test_SetWaterInfo_inexisting_value(
+def test_SetWaterInfo_invalid(
     command_values: dict[str, Any], error: type[Exception], error_message: str
 ) -> None:
     with pytest.raises(error, match=error_message):