fix: adding local_api disconnection

Author: humbertogontijo

roborock/api.py

@@ -106,6 +106,9 @@ def __init__(self, endpoint: str, device_localkey: dict[str, str]) -> None:
     def add_status_listener(self, callback: Callable[[str, str], None]):
         self._status_listeners.append(callback)
 
+    async def async_disconnect(self) -> Any:
+        raise NotImplementedError
+
     def _decode_msg(self, msg: bytes, local_key: str) -> dict[str, Any]:
         if msg[4:7] == "1.0".encode():
             msg = msg[4:]

roborock/local_api.py

@@ -4,7 +4,7 @@
 import logging
 import socket
 from asyncio import Lock
-from typing import Callable, Coroutine
+from typing import Callable, Coroutine, Any
 
 import async_timeout
 
@@ -37,6 +37,9 @@ async def async_connect(self):
             for listener in self.device_listener.values()
         ])
 
+    async def async_disconnect(self) -> Any:
+        await asyncio.gather(*[listener.disconnect() for listener in self.device_listener.values()])
+
     async def send_command(
             self, device_id: str, method: RoborockCommand, params: list = None
     ):
@@ -87,17 +90,23 @@ async def _main_coro(self):
                 except Exception as e:
                     _LOGGER.exception(e)
             except BrokenPipeError:
-                self.socket.close()
+                await self.disconnect()
 
     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)
                 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
                 self.loop.create_task(self._main_coro())
 
+    async def disconnect(self):
+        self.socket.close()
+        self.is_connected = False
+
     async def send_message(self, data: bytes):
         response = {}
         await self.connect()
@@ -110,5 +119,5 @@ async def send_message(self, data: bytes):
                 f"Timeout after {self.timeout} seconds waiting for response"
             ) from None
         except BrokenPipeError:
-            self.socket.close()
+            await self.disconnect()
         return response