fix: improving new protocols

Author: humbertogontijo

roborock/code_mappings.py

@@ -10,6 +10,10 @@
 class RoborockEnum(IntEnum):
     """Roborock Enum for codes with int values"""
 
+    @property
+    def name(self) -> str:
+        return super().name.lower()
+
     @classmethod
     def _missing_(cls: Type[RoborockEnum], key) -> RoborockEnum:
         if hasattr(cls, "unknown"):

roborock/local_api.py

@@ -2,16 +2,20 @@
 
 import asyncio
 import logging
-from asyncio import Lock, Transport
+from asyncio import Lock, TimerHandle, Transport
 from typing import Optional
 
 import async_timeout
 
-from . import DeviceData, RoborockBase
+from . import DeviceData
 from .api import COMMANDS_SECURED, QUEUE_TIMEOUT, RoborockClient
-from .exceptions import CommandVacuumError, RoborockConnectionException, RoborockException
+from .exceptions import (
+    CommandVacuumError,
+    RoborockConnectionException,
+    RoborockException,
+)
 from .protocol import MessageParser
-from .roborock_message import RoborockMessage
+from .roborock_message import RoborockMessage, RoborockMessageProtocol
 from .roborock_typing import CommandInfoMap, RoborockCommand
 from .util import get_running_loop_or_create_one
 
@@ -30,7 +34,7 @@ def __init__(self, device_data: DeviceData):
         self.remaining = b""
         self.transport: Transport | None = None
         self._mutex = Lock()
-        self.keep_alive_func()
+        self.keep_alive_task: TimerHandle | None = None
 
     def data_received(self, message):
         if self.remaining:
@@ -46,33 +50,34 @@ def connection_lost(self, exc: Optional[Exception]):
     def is_connected(self):
         return self.transport and self.transport.is_reading()
 
-    def keep_alive_func(self, _=None):
-        keep_alive_task = asyncio.gather(
-            asyncio.sleep(10),
-            self.ping(),
-        )
-        keep_alive_task.add_done_callback(self.keep_alive_func)
+    async def keep_alive_func(self, _=None):
+        await self.ping()
+        self.keep_alive_task = self.loop.call_later(10, lambda: self.keep_alive_func())
 
     async def async_connect(self) -> None:
         async with self._mutex:
             try:
-                if not self.is_connected():
+                is_connected = self.is_connected()
+                if not is_connected:
                     self.sync_disconnect()
                     async with async_timeout.timeout(QUEUE_TIMEOUT):
                         _LOGGER.info(f"Connecting to {self.host}")
                         self.transport, _ = await self.loop.create_connection(  # type: ignore
                             lambda: self, self.host, 58867
                         )
-                        await self.hello()
                         _LOGGER.info(f"Connected to {self.host}")
             except Exception as e:
-                _LOGGER.warning(f"Failed connecting to {self.host}: {e}")
                 raise RoborockConnectionException(f"Failed connecting to {self.host}") from e
+        if not is_connected:
+            await self.hello()
+            await self.keep_alive_func()
 
     def sync_disconnect(self) -> None:
         if self.transport and self.loop.is_running():
             _LOGGER.debug(f"Disconnecting from {self.host}")
             self.transport.close()
+        if self.keep_alive_task:
+            self.keep_alive_task.cancel()
 
     async def async_disconnect(self) -> None:
         async with self._mutex:
@@ -96,24 +101,28 @@ def build_roborock_message(self, method: RoborockCommand, params: Optional[list
         )
 
     async def hello(self):
-        request_id = 1
-        _LOGGER.debug(f"id={request_id} Requesting method hello with None")
-        try:
-            return await self.send_message(
-                RoborockMessage(protocol=0, payload=None, seq=request_id, version=b"1.0", random=22)
-            )
-        except Exception as e:
-            _LOGGER.error(e)
+        if self.is_connected():
+            request_id = 1
+            protocol = RoborockMessageProtocol.HELLO_REQUEST
+            _LOGGER.debug(f"id={request_id} Requesting protocol {protocol.name}")
+            try:
+                return await self.send_message(
+                    RoborockMessage(protocol=protocol, payload=None, seq=request_id, version=b"1.0", random=22)
+                )
+            except Exception as e:
+                _LOGGER.error(e)
 
     async def ping(self):
-        request_id = 2
-        _LOGGER.debug(f"id={request_id} Requesting method ping with None")
-        try:
-            return await self.send_message(
-                RoborockMessage(protocol=2, payload=None, seq=request_id, version=b"1.0", random=23)
-            )
-        except Exception as e:
-            _LOGGER.error(e)
+        if self.is_connected():
+            request_id = 2
+            protocol = RoborockMessageProtocol.PING_REQUEST
+            _LOGGER.debug(f"id={request_id} Requesting protocol {protocol.name}")
+            try:
+                return await self.send_message(
+                    RoborockMessage(protocol=protocol, payload=None, seq=request_id, version=b"1.0", random=23)
+                )
+            except Exception as e:
+                _LOGGER.error(e)
 
     async def send_command(self, method: RoborockCommand, params: Optional[list | dict] = None):
         roborock_message = self.build_roborock_message(method, params)
@@ -133,8 +142,12 @@ async def async_local_response(self, roborock_message: RoborockMessage):
         (response, err) = await self._async_response(request_id, response_protocol)
         if err:
             raise CommandVacuumError("", err) from err
-        method = roborock_message.get_method() if roborock_message.protocol != 2 else "ping"
-        _LOGGER.debug(f"id={request_id} Response from method {method}: {response}")
+        method = (
+            f"method {roborock_message.get_method()}"
+            if roborock_message.protocol == 4
+            else f"protocol {RoborockMessageProtocol(roborock_message.protocol).name}"
+        )
+        _LOGGER.debug(f"id={request_id} Response from {method}: {response}")
         return response
 
     def _send_msg_raw(self, data: bytes):
@@ -146,28 +159,22 @@ def _send_msg_raw(self, data: bytes):
             raise RoborockException(e) from e
 
     async def send_message(self, roborock_messages: list[RoborockMessage] | RoborockMessage):
-        try:
-            await self.validate_connection()
-            if isinstance(roborock_messages, RoborockMessage):
-                roborock_messages = [roborock_messages]
-            local_key = self.device_info.device.local_key
-            msg = MessageParser.build(roborock_messages, local_key=local_key)
-            # Send the command to the Roborock device
-            self._send_msg_raw(msg)
-
-            responses = await asyncio.gather(
-                *[self.async_local_response(roborock_message) for roborock_message in roborock_messages],
-                return_exceptions=True,
-            )
-            exception = next((response for response in responses if isinstance(response, BaseException)), None)
-            if exception:
-                raise exception
-            is_cached = next(
-                (response for response in responses if isinstance(response, RoborockBase) and response.is_cached), None
-            )
-            if is_cached:
-                await self.async_disconnect()
-            return responses
-        except Exception:
-            await self.async_disconnect()
-            raise
+        await self.validate_connection()
+        if isinstance(roborock_messages, RoborockMessage):
+            roborock_messages = [roborock_messages]
+        local_key = self.device_info.device.local_key
+        msg = MessageParser.build(roborock_messages, local_key=local_key)
+        # Send the command to the Roborock device
+        self._send_msg_raw(msg)
+
+        responses = await asyncio.gather(
+            *[self.async_local_response(roborock_message) for roborock_message in roborock_messages],
+            return_exceptions=True,
+        )
+        exception = next(
+            (response for response in responses if isinstance(response, BaseException)),
+            None,
+        )
+        if exception:
+            raise exception
+        return responses

roborock/roborock_message.py

@@ -6,6 +6,17 @@
 from dataclasses import dataclass
 from random import randint
 
+from roborock import RoborockEnum
+
+
+class RoborockMessageProtocol(RoborockEnum):
+    HELLO_REQUEST = 0
+    HELLO_RESPONSE = 1
+    PING_REQUEST = 2
+    PING_RESPONSE = 3
+    GENERAL_REQUEST = 4
+    GENERAL_RESPONSE = 5
+
 
 @dataclass
 class RoborockMessage: