Implement ResetLifeSpan

nanomad · nanomad · commit a4b9794fec14 · 2025-04-19T13:38:55.000+02:00
diff --git a/deebot_client/commands/__init__.py b/deebot_client/commands/__init__.py
@@ -11,14 +11,22 @@
     COMMANDS as JSON_COMMANDS,
     COMMANDS_WITH_MQTT_P2P_HANDLING as JSON_COMMANDS_WITH_MQTT_P2P_HANDLING,
 )
+from .xml import (
+    COMMANDS as XML_COMMANDS,
+    COMMANDS_WITH_MQTT_P2P_HANDLING as XML_COMMANDS_WITH_MQTT_P2P_HANDLING,
+)
 
 if TYPE_CHECKING:
     from deebot_client.command import Command, CommandMqttP2P
 
-COMMANDS: dict[DataType, dict[str, type[Command]]] = {DataType.JSON: JSON_COMMANDS}
+COMMANDS: dict[DataType, dict[str, type[Command]]] = {
+    DataType.JSON: JSON_COMMANDS,
+    DataType.XML: XML_COMMANDS,
+}
 
 COMMANDS_WITH_MQTT_P2P_HANDLING: dict[DataType, dict[str, type[CommandMqttP2P]]] = {
-    DataType.JSON: JSON_COMMANDS_WITH_MQTT_P2P_HANDLING
+    DataType.JSON: JSON_COMMANDS_WITH_MQTT_P2P_HANDLING,
+    DataType.XML: XML_COMMANDS_WITH_MQTT_P2P_HANDLING,
 }
 
 
diff --git a/deebot_client/commands/xml/__init__.py b/deebot_client/commands/xml/__init__.py
@@ -10,7 +10,7 @@
 from .charge_state import GetChargeState
 from .error import GetError
 from .fan_speed import GetFanSpeed
-from .life_span import GetLifeSpan
+from .life_span import GetLifeSpan, ResetLifeSpan
 from .play_sound import PlaySound
 from .pos import GetPos
 from .stats import GetCleanSum
@@ -27,14 +27,21 @@
     "GetLifeSpan",
     "GetPos",
     "PlaySound",
+    "ResetLifeSpan",
 ]
 
 # fmt: off
 # ordered by file asc
 _COMMANDS: list[type[XmlCommand]] = [
+    Charge,
+    GetChargeState,
+    GetCleanSum,
     GetError,
+    GetFanSpeed,
     GetLifeSpan,
+    GetPos,
     PlaySound,
+    ResetLifeSpan,
 ]
 # fmt: on
 
diff --git a/deebot_client/commands/xml/common.py b/deebot_client/commands/xml/common.py
@@ -3,17 +3,24 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, cast, override
 from xml.etree.ElementTree import Element, SubElement
 
 from defusedxml import ElementTree  # type: ignore[import-untyped]
 
-from deebot_client.command import Command, CommandWithMessageHandling, SetCommand
+from deebot_client.command import (
+    Command,
+    CommandMqttP2P,
+    CommandWithMessageHandling,
+    SetCommand,
+)
 from deebot_client.const import DataType
 from deebot_client.logging_filter import get_logger
 from deebot_client.message import HandlingResult, HandlingState, MessageStr
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from deebot_client.event_bus import EventBus
 
 _LOGGER = get_logger(__name__)
@@ -76,11 +83,46 @@ def _handle_xml(cls, _: EventBus, xml: Element) -> HandlingResult:
             return HandlingResult.success()
 
         _LOGGER.warning(
-            'Command "%s" was not successful. XML response: %s', cls.NAME, xml
+            'Command "%s" was not successful. XML response: %s',
+            cls.NAME,
+            ElementTree.tostring(xml, "unicode"),
         )
         return HandlingResult(HandlingState.FAILED)
 
 
+class XmlCommandMqttP2P(XmlCommand, CommandMqttP2P, ABC):
+    """Json base command for mqtt p2p channel."""
+
+    @classmethod
+    def create_from_mqtt(cls, payload: str | bytes | bytearray) -> CommandMqttP2P:
+        """Create a command from the mqtt data."""
+        xml = ElementTree.fromstring(payload)
+        return cls._create_from_mqtt(xml.attrib)
+
+    @override
+    def handle_mqtt_p2p(
+        self, event_bus: EventBus, response_payload: str | bytes | bytearray
+    ) -> None:
+        """Handle response received over the mqtt channel "p2p"."""
+        if isinstance(response_payload, bytearray):
+            data = bytes(response_payload).decode()
+        elif isinstance(response_payload, bytes):
+            data = response_payload.decode()
+        elif isinstance(response_payload, str):
+            data = response_payload
+        else:
+            msg = "Unsupported message data type {message_type}"  # type: ignore[unreachable]
+            raise TypeError(msg.format(message_type=type(response_payload)))
+
+        self._handle_mqtt_p2p(event_bus, data)
+
+    @abstractmethod
+    def _handle_mqtt_p2p(
+        self, event_bus: EventBus, response: dict[str, Any] | str
+    ) -> None:
+        """Handle response received over the mqtt channel "p2p"."""
+
+
 class XmlSetCommand(ExecuteCommand, SetCommand, ABC):
     """Xml base set command.
 
diff --git a/deebot_client/commands/xml/life_span.py b/deebot_client/commands/xml/life_span.py
@@ -2,12 +2,14 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any
 
+from deebot_client.command import InitParam
 from deebot_client.events import LifeSpan, LifeSpanEvent
-from deebot_client.message import HandlingResult
+from deebot_client.message import HandlingResult, HandlingState
 
-from .common import XmlCommandWithMessageHandling
+from .common import ExecuteCommand, XmlCommandMqttP2P, XmlCommandWithMessageHandling
 
 if TYPE_CHECKING:
     from xml.etree.ElementTree import Element
@@ -20,8 +22,11 @@ class GetLifeSpan(XmlCommandWithMessageHandling):
 
     NAME = "GetLifeSpan"
 
-    def __init__(self, life_span: LifeSpan) -> None:
-        super().__init__({"type": life_span.xml_value})
+    def __init__(self, life_span: LifeSpan | str) -> None:
+        xml_value = (
+            life_span.xml_value if isinstance(life_span, LifeSpan) else life_span
+        )
+        super().__init__({"type": xml_value})
 
     @classmethod
     def _handle_xml(cls, event_bus: EventBus, xml: Element) -> HandlingResult:
@@ -47,3 +52,24 @@ def _handle_xml(cls, event_bus: EventBus, xml: Element) -> HandlingResult:
             LifeSpanEvent(LifeSpan.from_xml(component_type), percent, left)
         )
         return HandlingResult.success()
+
+
+class ResetLifeSpan(ExecuteCommand, XmlCommandMqttP2P):
+    """ResetLifeSpan command."""
+
+    NAME = "ResetLifeSpan"
+    _mqtt_params = MappingProxyType({"type": InitParam(str, "life_span")})
+
+    def __init__(self, life_span: LifeSpan | str) -> None:
+        xml_value = (
+            life_span.xml_value if isinstance(life_span, LifeSpan) else life_span
+        )
+        super().__init__({"type": xml_value})
+
+    def _handle_mqtt_p2p(
+        self, event_bus: EventBus, response: dict[str, Any] | str
+    ) -> None:
+        """Handle response received over the mqtt channel "p2p"."""
+        result = self.handle(event_bus, response)
+        if result.state == HandlingState.SUCCESS:
+            event_bus.request_refresh(LifeSpanEvent)
diff --git a/tests/commands/xml/__init__.py b/tests/commands/xml/__init__.py
@@ -1,6 +1,61 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any, cast
+from xml.etree.ElementTree import Element, SubElement
+
+from defusedxml import ElementTree  # type: ignore[import-untyped]
+from testfixtures import LogCapture
+
+from deebot_client.command import CommandResult
+from deebot_client.message import HandlingState
+from tests.commands import assert_command
+
+if TYPE_CHECKING:
+    from deebot_client.commands.xml.common import ExecuteCommand
+
+
+def get_success_body(
+    extra_attrs: dict[str, Any] | None = None, sub_element_name: str | None = None
+) -> str:
+    element = ctl_element = Element("ctl")
+    element.set("ret", "ok")
+
+    if extra_attrs is not None and len(extra_attrs) > 0:
+        if sub_element_name is not None:
+            element = SubElement(element, sub_element_name.lower())
+
+        if isinstance(extra_attrs, dict):
+            for key, value in extra_attrs.items():
+                element.set(key, value)
+
+    return cast("str", ElementTree.tostring(ctl_element, "unicode"))
+
+
+async def assert_execute_command(
+    command: ExecuteCommand, args: dict[str, Any] | list[Any] | None
+) -> None:
+    assert command.NAME != "invalid"
+    assert command._args == args
+
+    # success
+    xml = get_request_xml(get_success_body())
+    await assert_command(command, xml, None)
+
+    # failed
+    with LogCapture() as log:
+        body = '<ctl ret="error" />'
+        json = get_request_xml(body)
+        await assert_command(
+            command, json, None, command_result=CommandResult(HandlingState.FAILED)
+        )
+
+        log.check_present(
+            (
+                "deebot_client.commands.xml.common",
+                "WARNING",
+                f'Command "{command.NAME}" was not successful. XML response: {body}',
+            )
+        )
 
 
 def get_request_xml(data: str | None) -> dict[str, Any]:
diff --git a/tests/commands/xml/test_common.py b/tests/commands/xml/test_common.py
@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from unittest.mock import Mock, patch
+
+import pytest
+
+from deebot_client.commands.xml.common import XmlCommandMqttP2P
+from deebot_client.event_bus import EventBus
+
+
+@pytest.mark.parametrize(
+    ("payload", "decoded_payload"),
+    [
+        (bytearray(bytes("test", "utf-8")), "test"),
+        (bytes("test", "utf-8"), "test"),
+        ("test", "test"),
+    ],
+    ids=["bytearray", "bytes", "str"],
+)
+@patch.multiple(XmlCommandMqttP2P, __abstractmethods__=set())
+def test_XmlCommandMqttP2P_decoding(
+    payload: bytearray | bytes | str, decoded_payload: str
+) -> None:
+    command = XmlCommandMqttP2P()  # type: ignore[abstract]
+    event_bus = Mock(spec_set=EventBus)
+    with patch.object(command, "_handle_mqtt_p2p", return_value=None) as mqtt_handler:
+        command.handle_mqtt_p2p(event_bus, payload)
+
+    mqtt_handler.assert_called_once_with(event_bus, decoded_payload)
+
+
+@patch.multiple(XmlCommandMqttP2P, __abstractmethods__=set())
+def test_XmlCommandMqttP2P_invalid_decoding() -> None:
+    command = XmlCommandMqttP2P()  # type: ignore[abstract]
+    event_bus = Mock(spec_set=EventBus)
+    with (
+        patch.object(command, "_handle_mqtt_p2p", return_value=None),
+        pytest.raises(TypeError),
+    ):
+        command.handle_mqtt_p2p(event_bus, {})  # type: ignore[arg-type]
diff --git a/tests/commands/xml/test_life_span.py b/tests/commands/xml/test_life_span.py
@@ -1,20 +1,33 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+from unittest.mock import patch
+
 import pytest
 
 from deebot_client.command import CommandResult
-from deebot_client.commands.xml import GetLifeSpan
+from deebot_client.commands.xml import GetLifeSpan, ResetLifeSpan
 from deebot_client.events import LifeSpan, LifeSpanEvent
 from deebot_client.message import HandlingState
 from tests.commands import assert_command
 
-from . import get_request_xml
+from . import assert_execute_command, get_request_xml, get_success_body
+
+if TYPE_CHECKING:
+    from deebot_client.event_bus import EventBus
 
 
 @pytest.mark.parametrize(
     ("component_type", "lifespan_type", "left", "total", "expected_event"),
     [
         ("Brush", LifeSpan.BRUSH, 50, 100, LifeSpanEvent(LifeSpan.BRUSH, 50, 50)),
+        (
+            "Brush",
+            LifeSpan.BRUSH.xml_value,
+            50,
+            100,
+            LifeSpanEvent(LifeSpan.BRUSH, 50, 50),
+        ),
         (
             "DustCaseHeap",
             LifeSpan.DUST_CASE_HEAP,
@@ -33,15 +46,15 @@
 )
 async def test_get_life_span(
     component_type: str,
-    lifespan_type: LifeSpan,
+    lifespan_type: LifeSpan | str,
     left: int,
     total: int,
     expected_event: LifeSpanEvent,
 ) -> None:
-    json = get_request_xml(
+    xml = get_request_xml(
         f"<ctl ret='ok' type='{component_type}' left='{left}' total='{total}'/>"
     )
-    await assert_command(GetLifeSpan(lifespan_type), json, expected_event)
+    await assert_command(GetLifeSpan(lifespan_type), xml, expected_event)
 
 
 @pytest.mark.parametrize(
@@ -57,3 +70,30 @@ async def test_get_life_span_error(xml: str) -> None:
         None,
         command_result=CommandResult(HandlingState.ANALYSE_LOGGED),
     )
+
+
+@pytest.mark.parametrize(
+    ("command", "args"),
+    [
+        (ResetLifeSpan(LifeSpan.FILTER), {"type": LifeSpan.FILTER.xml_value}),
+        (ResetLifeSpan(LifeSpan.FILTER.xml_value), {"type": LifeSpan.FILTER.xml_value}),
+        (
+            ResetLifeSpan.create_from_mqtt(b'<ctl type="Brush" />'),
+            {"type": LifeSpan.BRUSH.xml_value},
+        ),
+    ],
+)
+async def test_ResetLifeSpan(command: ResetLifeSpan, args: dict[str, str]) -> None:
+    await assert_execute_command(command, args)
+
+
+def test_ResetLifeSpan_invokes_refresh(event_bus: EventBus) -> None:
+    command = ResetLifeSpan(LifeSpan.FILTER)
+    success_response = get_success_body()
+
+    with patch.object(
+        event_bus, "request_refresh", return_value=None
+    ) as mock_request_refresh:
+        command.handle_mqtt_p2p(event_bus, success_response)
+
+    mock_request_refresh.assert_called_with(LifeSpanEvent)
