fix: improving local integration

humbertogontijo · humbertogontijo · commit 7657617901d8 · 2023-04-21T23:10:16.000-03:00
diff --git a/roborock/api.py b/roborock/api.py
@@ -194,10 +194,7 @@ async def get_clean_summary(self, device_id: str) -> CleanSummary | None:
             elif isinstance(clean_summary, list):
                 clean_time, clean_area, clean_count, records = unpack_list(clean_summary, 4)
                 return CleanSummary(
-                    clean_time=clean_time,
-                    clean_area=clean_area,
-                    clean_count=clean_count,
-                    records=records
+                    clean_time=clean_time, clean_area=clean_area, clean_count=clean_count, records=records
                 )
             elif isinstance(clean_summary, int):
                 return CleanSummary(clean_time=clean_summary)
@@ -268,7 +265,9 @@ async def get_dock_summary(self, device_id: str, dock_type: RoborockEnum) -> Rob
                     self.get_wash_towel_mode(device_id),
                     self.get_smart_wash_params(device_id),
                 ]
-            [dust_collection_mode, wash_towel_mode, smart_wash_params] = unpack_list(list(await asyncio.gather(*commands)), 3)
+            [dust_collection_mode, wash_towel_mode, smart_wash_params] = unpack_list(
+                list(await asyncio.gather(*commands)), 3
+            )
 
             return RoborockDockSummary(dust_collection_mode, wash_towel_mode, smart_wash_params)
         except RoborockTimeout as e:
diff --git a/roborock/cloud_api.py b/roborock/cloud_api.py
@@ -5,7 +5,7 @@
 import threading
 import uuid
 from asyncio import Lock
-from typing import Any, Mapping, Optional
+from typing import Mapping, Optional
 from urllib.parse import urlparse
 
 import paho.mqtt.client as mqtt
@@ -137,34 +137,30 @@ def sync_disconnect(self) -> bool:
         return rc == mqtt.MQTT_ERR_SUCCESS
 
     def sync_connect(self) -> bool:
-        rc = mqtt.MQTT_ERR_AGAIN
         self.sync_start_loop()
         if not self.is_connected():
             if self._mqtt_port is None or self._mqtt_host is None:
                 raise RoborockException("Mqtt information was not entered. Cannot connect.")
             _LOGGER.info("Connecting to mqtt")
-            rc = super().connect(host=self._mqtt_host, port=self._mqtt_port, keepalive=MQTT_KEEPALIVE)
-            if rc != mqtt.MQTT_ERR_SUCCESS:
-                raise RoborockException(f"Failed to connect (rc:{rc})")
-        return rc == mqtt.MQTT_ERR_SUCCESS
+            super().connect_async(host=self._mqtt_host, port=self._mqtt_port, keepalive=MQTT_KEEPALIVE)
+            return True
+        return False
 
-    async def async_disconnect(self) -> Any:
+    async def async_disconnect(self) -> None:
         async with self._mutex:
             disconnecting = self.sync_disconnect()
             if disconnecting:
-                (response, err) = await self._async_response(DISCONNECT_REQUEST_ID)
+                (_, err) = await self._async_response(DISCONNECT_REQUEST_ID)
                 if err:
                     raise RoborockException(err) from err
-                return response
 
-    async def async_connect(self) -> Any:
+    async def async_connect(self) -> None:
         async with self._mutex:
             connecting = self.sync_connect()
             if connecting:
-                (response, err) = await self._async_response(CONNECT_REQUEST_ID)
+                (_, err) = await self._async_response(CONNECT_REQUEST_ID)
                 if err:
                     raise RoborockException(err) from err
-                return response
 
     async def validate_connection(self) -> None:
         await self.async_connect()
@@ -194,4 +190,8 @@ async def send_command(self, device_id: str, method: RoborockCommand, params: Op
         return response
 
     async def get_map_v1(self, device_id):
-        return await self.send_command(device_id, RoborockCommand.GET_MAP_V1)
+        try:
+            return await self.send_command(device_id, RoborockCommand.GET_MAP_V1)
+        except RoborockException as e:
+            _LOGGER.error(e)
+        return None
diff --git a/roborock/containers.py b/roborock/containers.py
@@ -35,6 +35,8 @@ def decamelize(s: str):
 
 
 def decamelize_obj(d: dict | list, ignore_keys: list[str]):
+    if isinstance(d, RoborockBase):
+        d = d.as_dict()
     if isinstance(d, list):
         return [decamelize_obj(i, ignore_keys) if isinstance(i, (dict, list)) else i for i in d]
     return {
diff --git a/roborock/local_api.py b/roborock/local_api.py
@@ -3,14 +3,14 @@
 import asyncio
 import logging
 import socket
-from asyncio import Lock
+from asyncio import Transport
 from typing import Callable, Mapping, Optional
 
 import async_timeout
 
-from .api import SPECIAL_COMMANDS, RoborockClient
+from .api import SPECIAL_COMMANDS, RoborockClient, QUEUE_TIMEOUT
 from .containers import RoborockLocalDeviceInfo
-from .exceptions import CommandVacuumError, RoborockConnectionException, RoborockException, RoborockTimeout
+from .exceptions import CommandVacuumError, RoborockConnectionException, RoborockException
 from .roborock_message import RoborockMessage, RoborockParser
 from .typing import CommandInfoMap, RoborockCommand
 from .util import get_running_loop_or_create_one
@@ -30,15 +30,15 @@ def __init__(self, devices_info: Mapping[str, RoborockLocalDeviceInfo]):
             )
             for device_id, device_info in devices_info.items()
         }
-        self._mutex = Lock()
         self._batch_structs: list[RoborockMessage] = []
         self._executing = False
 
-    async def async_connect(self):
+    async def async_connect(self) -> None:
         await asyncio.gather(*[listener.connect() for listener in self.device_listener.values()])
 
     async def async_disconnect(self) -> None:
-        await asyncio.gather(*[listener.disconnect() for listener in self.device_listener.values()])
+        for listener in self.device_listener.values():
+            listener.disconnect()
 
     def build_roborock_message(self, method: RoborockCommand, params: Optional[list] = None) -> RoborockMessage:
         secured = True if method in SPECIAL_COMMANDS else False
@@ -103,72 +103,56 @@ def is_closed(self):
         return self._closed
 
 
-class RoborockSocketListener:
+class RoborockSocketListener(asyncio.Protocol):
     roborock_port = 58867
 
     def __init__(
         self,
         ip: str,
         local_key: str,
         on_message: Callable[[list[RoborockMessage]], None],
-        timeout: float | int = 4,
+        timeout: float | int = QUEUE_TIMEOUT,
     ):
         self.ip = ip
         self.local_key = local_key
-        self.socket = RoborockSocket(socket.AF_INET, socket.SOCK_STREAM)
-        self.socket.setblocking(False)
         self.loop = get_running_loop_or_create_one()
         self.on_message = on_message
         self.timeout = timeout
-        self.is_connected = False
-        self._mutex = Lock()
         self.remaining = b""
+        self.transport: Transport | None = None
 
-    async def _main_coro(self):
-        while not self.socket.is_closed:
-            try:
-                message = await self.loop.sock_recv(self.socket, 4096)
-                try:
-                    if self.remaining:
-                        message = self.remaining + message
-                        self.remaining = b""
-                    (parser_msg, remaining) = RoborockParser.decode(message, self.local_key)
-                    self.remaining = remaining
-                    self.on_message(parser_msg)
-                except Exception as e:
-                    _LOGGER.exception(e)
-            except BrokenPipeError as e:
-                _LOGGER.exception(e)
-                await self.disconnect()
+    def data_received(self, message):
+        if self.remaining:
+            message = self.remaining + message
+            self.remaining = b""
+        (parser_msg, remaining) = RoborockParser.decode(message, self.local_key)
+        self.remaining = remaining
+        self.on_message(parser_msg)
+
+    def connection_lost(self, exc):
+        print("The server closed the connection")
+
+    def is_connected(self):
+        return self.transport and self.transport.is_reading()
 
     async def connect(self):
-        async with self._mutex:
-            if not self.is_connected or self.socket.is_closed:
-                self.socket = RoborockSocket(socket.AF_INET, socket.SOCK_STREAM)
-                self.socket.setblocking(False)
-                try:
-                    async with async_timeout.timeout(self.timeout):
-                        _LOGGER.info(f"Connecting to {self.ip}")
-                        await self.loop.sock_connect(self.socket, (self.ip, 58867))
-                        self.is_connected = True
-                except Exception as e:
-                    await self.disconnect()
-                    raise RoborockConnectionException(f"Failed connecting to {self.ip}") from e
-                self.loop.create_task(self._main_coro())
-
-    async def disconnect(self):
-        self.socket.close()
-        self.is_connected = False
+        try:
+            if not self.is_connected():
+                async with async_timeout.timeout(self.timeout):
+                    _LOGGER.info(f"Connecting to {self.ip}")
+                    self.transport, _ = await self.loop.create_connection(lambda: self, self.ip, 58867)
+        except Exception as e:
+            raise RoborockConnectionException(f"Failed connecting to {self.ip}") from e
+
+    def disconnect(self):
+        if self.transport:
+            self.transport.close()
 
     async def send_message(self, data: bytes) -> None:
         await self.connect()
         try:
-            async with self._mutex:
-                async with async_timeout.timeout(self.timeout):
-                    await self.loop.sock_sendall(self.socket, data)
-        except (asyncio.TimeoutError, asyncio.CancelledError):
-            await self.disconnect()
-            raise RoborockTimeout(f"Timeout after {self.timeout} seconds waiting for response") from None
-        except BrokenPipeError as e:
-            await self.disconnect()
+            if not self.transport:
+                raise RoborockException("Can not send message without connection")
+            self.transport.write(data)
+        except Exception as e:
             raise RoborockException(e) from e
diff --git a/roborock/roborock_future.py b/roborock/roborock_future.py
@@ -18,5 +18,8 @@ def resolve(self, item: tuple[Any, VacuumError | None]) -> None:
         self.loop.call_soon_threadsafe(self.fut.set_result, item)
 
     async def async_get(self, timeout: float | int) -> tuple[Any, VacuumError | None]:
-        async with async_timeout.timeout(timeout):
-            return await self.fut
+        try:
+            async with async_timeout.timeout(timeout):
+                return await self.fut
+        finally:
+            self.fut.cancel()
diff --git a/roborock/util.py b/roborock/util.py
@@ -7,9 +7,11 @@
 
 T = TypeVar("T")
 
+
 def unpack_list(value: list[T], size: int) -> list[T | None]:
     return (value + [None] * size)[:size]
 
+
 def get_running_loop_or_create_one() -> AbstractEventLoop:
     try:
         loop = asyncio.get_event_loop()
