Add first batch of XML Water related commands and messages

Author: nanomad

deebot_client/commands/xml/__init__.py

@@ -18,6 +18,7 @@
 from .play_sound import PlaySound
 from .pos import GetChargerPos, GetPos
 from .stats import GetCleanSum
+from .water_info import GetWaterBoxInfo, GetWaterPermeability, SetWaterPermeability
 
 if TYPE_CHECKING:
     from .common import XmlCommand
@@ -40,10 +41,13 @@
     "GetMapSt",
     "GetPos",
     "GetTrM",
+    "GetWaterBoxInfo",
+    "GetWaterPermeability",
     "PlaySound",
     "PullM",
     "PullMP",
     "SetCleanSpeed",
+    "SetWaterPermeability",
 ]
 
 # fmt: off
@@ -81,6 +85,10 @@
     GetPos,
 
     GetCleanSum,
+
+    GetWaterPermeability,
+    GetWaterBoxInfo,
+    SetWaterPermeability,
 ]
 # fmt: on
 

deebot_client/commands/xml/water_info.py

@@ -0,0 +1,71 @@
+"""WaterBox command module."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from deebot_client.commands.xml.common import (
+    ExecuteCommand,
+    XmlCommandWithMessageHandling,
+)
+from deebot_client.events.water_info import (
+    WaterAmount,
+    WaterAmountEvent,
+)
+from deebot_client.message import HandlingResult
+from deebot_client.messages.xml import WaterBoxInfo
+from deebot_client.util import get_enum
+
+if TYPE_CHECKING:
+    from xml.etree.ElementTree import Element
+
+    from deebot_client.event_bus import EventBus
+
+
+class GetWaterPermeability(XmlCommandWithMessageHandling):
+    """GetWaterPermeability command."""
+
+    NAME = "GetWaterPermeability"
+
+    @classmethod
+    def _handle_xml(cls, event_bus: EventBus, xml: Element) -> HandlingResult:
+        """Handle xml message and notify the correct event subscribers.
+
+        :return: A message response.
+        """
+        if xml.attrib.get("ret") != "ok" or not (value := xml.attrib.get("v")):
+            return HandlingResult.analyse()
+
+        if value.isdecimal() and (value_int := int(value)) >= 0:
+            event_bus.notify(WaterAmountEvent(WaterAmount(value_int)))
+            return HandlingResult.success()
+
+        return HandlingResult.analyse()
+
+
+class SetWaterPermeability(ExecuteCommand):
+    """SetWaterPermeability command."""
+
+    NAME = "SetWaterPermeability"
+
+    def __init__(self, amount: WaterAmount | str) -> None:
+        if isinstance(amount, str):
+            amount = get_enum(WaterAmount, amount)
+        super().__init__({"v": str(amount.value)})
+
+
+class GetWaterBoxInfo(XmlCommandWithMessageHandling, WaterBoxInfo):
+    """GetWaterBoxInfo command."""
+
+    NAME = "GetWaterBoxInfo"
+
+    @classmethod
+    def _handle_xml(cls, event_bus: EventBus, xml: Element) -> HandlingResult:
+        """Handle xml message and notify the correct event subscribers.
+
+        :return: A message response.
+        """
+        if xml.attrib.get("ret") != "ok":
+            return HandlingResult.analyse()
+
+        return cls._parse_xml(event_bus, xml)

deebot_client/messages/xml/__init__.py

@@ -7,21 +7,24 @@
 from deebot_client.messages.xml.battery import BatteryInfo
 from deebot_client.messages.xml.charge_state import ChargeState
 from deebot_client.messages.xml.pos import Pos
+from deebot_client.messages.xml.water_info import WaterBoxInfo
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
     from deebot_client.message import Message
 
-__all__: Sequence[str] = ["BatteryInfo", "ChargeState", "Pos"]
+__all__: Sequence[str] = ["BatteryInfo", "ChargeState", "Pos", "WaterBoxInfo"]
 # fmt: off
 # ordered by file asc
 _MESSAGES: list[type[Message]] = [
     BatteryInfo,
 
     ChargeState,
 
-    Pos
+    Pos,
+
+    WaterBoxInfo
 ]
 # fmt: on
 

deebot_client/messages/xml/water_info.py

@@ -0,0 +1,40 @@
+"""Water messages."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from deebot_client.events.water_info import MopAttachedEvent
+from deebot_client.message import HandlingResult
+from deebot_client.messages.xml.common import XmlMessage
+
+if TYPE_CHECKING:
+    from xml.etree.ElementTree import Element
+
+    from deebot_client.event_bus import EventBus
+
+
+class WaterBoxInfo(XmlMessage):
+    """WaterBoxInfo message."""
+
+    NAME = "WaterBoxInfo"
+
+    @classmethod
+    def _handle_xml(cls, event_bus: EventBus, xml: Element) -> HandlingResult:
+        """Handle xml message and notify the correct event subscribers.
+
+        :return: A message response.
+        """
+        return cls._parse_xml(event_bus, xml)
+
+    @classmethod
+    def _parse_xml(cls, event_bus: EventBus, xml: Element) -> HandlingResult:
+        """Handle xml message and notify the correct event subscribers.
+
+        :return: A message response.
+        """
+        if (on := xml.attrib.get("on")) is None:
+            return HandlingResult.analyse()
+
+        event_bus.notify(MopAttachedEvent(on != "0"))
+        return HandlingResult.success()

tests/commands/xml/test_water_info.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+from deebot_client.command import CommandResult
+from deebot_client.commands.xml import GetWaterBoxInfo
+from deebot_client.events.water_info import MopAttachedEvent
+from deebot_client.message import HandlingState
+from tests.commands import assert_command
+
+from . import get_request_xml
+
+if TYPE_CHECKING:
+    from deebot_client.events.base import Event
+
+
+@pytest.mark.parametrize(
+    ("state", "expected_event"),
+    [
+        (1, MopAttachedEvent(True)),
+        (0, MopAttachedEvent(False)),
+    ],
+    ids=["mop_attached", "mop_not_attached"],
+)
+async def test_get_water_box_info(state: int, expected_event: Event) -> None:
+    xml_message = get_request_xml(f"<ctl ret='ok' on='{state}' />")
+    await assert_command(GetWaterBoxInfo(), xml_message, expected_event)
+
+
+@pytest.mark.parametrize(
+    "payload",
+    [
+        '<ctl ret="error"/>',
+        '<ctl ret="ok"></ctl>',
+    ],
+    ids=[
+        "error",
+        "no_state",
+    ],
+)
+async def test_get_water_box_info_error(payload: str) -> None:
+    xml_message = get_request_xml(payload)
+    await assert_command(
+        GetWaterBoxInfo(),
+        xml_message,
+        None,
+        command_result=CommandResult(HandlingState.ANALYSE_LOGGED),
+    )