From 725fb6bb98763c58fabf560bb7146ee0d9fb0505 Mon Sep 17 00:00:00 2001
From: Wentao Wu <wuwentao@canaan-creative.com>
Date: Wed, 16 Oct 2024 09:05:14 +0000
Subject: [PATCH] feat(ed): add new query01 msg for ed device

---
 midealocal/devices/ed/__init__.py   | 15 ++++++++++++---
 midealocal/devices/ed/message.py    | 26 +++++++++++++++++++++++++-
 tests/devices/ed/device_ed_test.py  |  5 +++--
 tests/devices/ed/message_ed_test.py | 16 ++++++++++++++--
 4 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/midealocal/devices/ed/__init__.py b/midealocal/devices/ed/__init__.py
index 92b62ddd..4c0c2dcf 100644
--- a/midealocal/devices/ed/__init__.py
+++ b/midealocal/devices/ed/__init__.py
@@ -8,7 +8,13 @@
 from midealocal.device import MideaDevice
 from midealocal.message import ListTypes
 
-from .message import MessageEDResponse, MessageNewSet, MessageOldSet, MessageQuery
+from .message import (
+    MessageEDResponse,
+    MessageNewSet,
+    MessageOldSet,
+    MessageQuery,
+    MessageQuery01,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -77,9 +83,12 @@ def _use_new_set(self) -> bool:
         # if (self.sub_type > 342 or self.sub_type == 340) else False
         return True
 
-    def build_query(self) -> list[MessageQuery]:
+    def build_query(self) -> list[MessageQuery | MessageQuery01]:
         """Midea ED device build query."""
-        return [MessageQuery(self._message_protocol_version, self._device_class)]
+        return [
+            MessageQuery(self._message_protocol_version),
+            MessageQuery01(self._message_protocol_version),
+        ]
 
     def process_message(self, msg: bytes) -> dict[str, Any]:
         """Midea ED device process message."""
diff --git a/midealocal/devices/ed/message.py b/midealocal/devices/ed/message.py
index ecad02e6..2d407065 100644
--- a/midealocal/devices/ed/message.py
+++ b/midealocal/devices/ed/message.py
@@ -64,7 +64,11 @@ def _body(self) -> bytearray:
 class MessageQuery(MessageEDBase):
     """ED message query."""
 
-    def __init__(self, protocol_version: int, body_type: ListTypes) -> None:
+    def __init__(
+        self,
+        protocol_version: int,
+        body_type: ListTypes = ListTypes.X00,
+    ) -> None:
         """Initialize ED message query."""
         super().__init__(
             protocol_version=protocol_version,
@@ -77,6 +81,26 @@ def _body(self) -> bytearray:
         return bytearray([0x01])
 
 
+class MessageQuery01(MessageEDBase):
+    """ED message query01."""
+
+    def __init__(
+        self,
+        protocol_version: int,
+        body_type: ListTypes = ListTypes.X01,
+    ) -> None:
+        """Initialize ED message query01."""
+        super().__init__(
+            protocol_version=protocol_version,
+            message_type=MessageType.query,
+            body_type=body_type,
+        )
+
+    @property
+    def _body(self) -> bytearray:
+        return bytearray([0x01])
+
+
 class MessageNewSet(MessageEDBase):
     """ED message new set."""
 
diff --git a/tests/devices/ed/device_ed_test.py b/tests/devices/ed/device_ed_test.py
index d29d7129..ca41f6fe 100644
--- a/tests/devices/ed/device_ed_test.py
+++ b/tests/devices/ed/device_ed_test.py
@@ -6,7 +6,7 @@
 
 from midealocal.const import ProtocolVersion
 from midealocal.devices.ed import DeviceAttributes, MideaEDDevice
-from midealocal.devices.ed.message import MessageQuery
+from midealocal.devices.ed.message import MessageQuery, MessageQuery01
 
 
 class TestMideaEDDevice:
@@ -89,8 +89,9 @@ def test_process_message(self) -> None:
     def test_build_query(self) -> None:
         """Test build query."""
         queries = self.device.build_query()
-        assert len(queries) == 1
+        assert len(queries) == 2
         assert isinstance(queries[0], MessageQuery)
+        assert isinstance(queries[1], MessageQuery01)
 
     def test_set_attribute(self) -> None:
         """Test set attribute."""
diff --git a/tests/devices/ed/message_ed_test.py b/tests/devices/ed/message_ed_test.py
index 3e280f62..81584953 100644
--- a/tests/devices/ed/message_ed_test.py
+++ b/tests/devices/ed/message_ed_test.py
@@ -14,6 +14,7 @@
     MessageEDResponse,
     MessageNewSet,
     MessageQuery,
+    MessageQuery01,
 )
 from midealocal.message import ListTypes, MessageType
 
@@ -39,9 +40,20 @@ def test_query_body(self) -> None:
         """Test query body."""
         query = MessageQuery(
             protocol_version=ProtocolVersion.V1,
-            body_type=ListTypes.X02,
         )
-        expected_body = bytearray([0x02, 0x01])
+        expected_body = bytearray([0x00, 0x01])
+        assert query.body == expected_body
+
+
+class TestMessageQuery01:
+    """Test Message Query01."""
+
+    def test_query_body(self) -> None:
+        """Test query body."""
+        query = MessageQuery01(
+            protocol_version=ProtocolVersion.V1,
+        )
+        expected_body = bytearray([0x01, 0x01])
         assert query.body == expected_body