Adding acceptable timeouts countdown

custom_components/roborock/__init__.py

@@ -96,7 +96,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     client = RoborockMqttClient(user_data, device_map)
     coordinator = RoborockDataUpdateCoordinator(hass, client, translation)
 
-    await coordinator.init()
     await coordinator.async_refresh()
 
     if not coordinator.last_update_success:
@@ -120,6 +119,8 @@ class RoborockDataUpdateCoordinator(
 ):
     """Class to manage fetching data from the API."""
 
+    ACCEPTABLE_NUMBER_OF_TIMEOUTS = 3
+
     def __init__(
             self, hass: HomeAssistant, client: RoborockMqttClient, translation: dict
     ) -> None:
@@ -130,30 +131,49 @@ def __init__(
         self._devices_prop: dict[str, RoborockDeviceProp] = {}
         self.translation = translation
         self.devices_maps: dict[str, MultiMapsList] = {}
+        self.retries = int(self.ACCEPTABLE_NUMBER_OF_TIMEOUTS)
+        self._timeout_countdown = int(self.ACCEPTABLE_NUMBER_OF_TIMEOUTS)
 
     def release(self):
         """Disconnect from API."""
         self.api.release()
 
-    async def init(self):
+    async def _get_device_multi_maps_list(self, device_id: str):
         """Get multi maps list."""
-        for device_id, _ in self.api.device_map.items():
-            multi_maps_list = await self.api.get_multi_maps_list(device_id)
+        multi_maps_list = await self.api.get_multi_maps_list(device_id)
+        if multi_maps_list:
             self.devices_maps[device_id] = multi_maps_list
 
+    async def _get_device_prop(self, device_id: str):
+        """Get device properties."""
+        device_prop = await self.api.get_prop(device_id)
+        if device_id in self._devices_prop:
+            self._devices_prop[device_id].update(device_prop)
+        else:
+            self._devices_prop[device_id] = device_prop
+
     async def _async_update_data(self):
         """Update data via library."""
         try:
+            funcs = []
             for device_id, _ in self.api.device_map.items():
-                device_prop = await self.api.get_prop(device_id)
-                if device_prop:
-                    if device_id in self._devices_prop:
-                        self._devices_prop[device_id].update(device_prop)
-                    else:
-                        self._devices_prop[device_id] = device_prop
-            return self._devices_prop
+                if not self.devices_maps.get(device_id):
+                    funcs.append(self._get_device_multi_maps_list(device_id))
+                funcs.append(self._get_device_prop(device_id))
+            await asyncio.gather(*funcs)
+            self._timeout_countdown = int(self.ACCEPTABLE_NUMBER_OF_TIMEOUTS)
         except (RoborockTimeout, RoborockException) as ex:
-            raise UpdateFailed(ex) from ex
+            if self._timeout_countdown > 0:
+                _LOGGER.debug("Timeout updating coordinator. Acceptable timeouts countdown = %s", self._timeout_countdown)
+                self._timeout_countdown -= 1
+            else:
+                raise UpdateFailed(ex) from ex
+        if self._devices_prop:
+            return self._devices_prop
+        # Only for the first attempt
+        if self.retries > 0:
+            self.retries -= 1
+            return self._async_update_data()
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:

custom_components/roborock/api/api.py

@@ -32,7 +32,6 @@
 _LOGGER = logging.getLogger(__name__)
 QUEUE_TIMEOUT = 4
 MQTT_KEEPALIVE = 60
-SESSION_EXPIRY_INTERVAL = 2 * 60
 
 
 def md5hex(message: str):
@@ -101,8 +100,7 @@ def __init__(self, user_data: UserData, device_map: dict[str, RoborockDeviceInfo
         self._nonce = secrets.token_bytes(16)
         self._waiting_queue: dict[int, RoborockQueue] = {}
         self._user_data = user_data
-        self._first_connection = True
-        self._last_message_timestamp = time.time()
+        self._should_connect = True
         self._mutex = Lock()
         super().__init__(client_id=self._hashed_user, protocol=mqtt.MQTTv5)
         if self._mqtt_ssl:
@@ -118,7 +116,6 @@ def __del__(self):
 
     @run_in_executor()
     async def on_connect(self, _client: mqtt.Client, _, __, rc, ___=None):
-        self._last_message_timestamp = time.time()
         connection_queue = self._waiting_queue[0]
         if rc != 0:
             await connection_queue.async_put((None, RoborockException("Failed to connect.")), timeout=QUEUE_TIMEOUT)
@@ -133,7 +130,6 @@ async def on_connect(self, _client: mqtt.Client, _, __, rc, ___=None):
     @run_in_executor()
     async def on_message(self, _client, _, msg, __=None):
         try:
-            self._last_message_timestamp = time.time()
             device_id = msg.topic.split("/").pop()
             data = self._decode_msg(msg.payload, self.device_map[device_id].device)
             protocol = data.get("protocol")
@@ -215,15 +211,16 @@ async def connect(self, **kwargs):
             if err:
                 raise RoborockException(err) from err
         except TimeoutError as ex:
+            self._should_connect = True
             raise RoborockTimeout(f"Timeout after {QUEUE_TIMEOUT} seconds waiting for mqtt connection") from ex
         finally:
             del self._waiting_queue[0]
 
     async def validate_connection(self):
         async with self._mutex:
-            if self._first_connection or time.time() - self._last_message_timestamp > SESSION_EXPIRY_INTERVAL:
+            if self._should_connect:
                 await self.async_reconnect()
-                self._first_connection = False
+                self._should_connect = False
 
     def _decode_msg(self, msg, device: HomeDataDevice):
         if msg[0:3] != "1.0".encode():
@@ -345,7 +342,6 @@ async def get_multi_maps_list(self, device_id):
             return MultiMapsList(multi_maps_list)
 
 
-
 class RoborockClient:
     def __init__(self, username: str, base_url=None) -> None:
         """Sample API Client."""

custom_components/roborock/config_flow.py

@@ -140,7 +140,7 @@ async def _login(self, code) -> UserData:
         except Exception as ex:
             _LOGGER.exception(ex)
             self._errors["base"] = "auth"
-            return None
+            return
 
 
 PERCENT_SCHEMA = vol.All(vol.Coerce(float), vol.Range(min=0, max=100))

custom_components/roborock/manifest.json

@@ -12,7 +12,7 @@
     "pyasn1",
     "pillow"
   ],
-  "version": "0.0.25",
+  "version": "0.0.26",
   "iot_class": "cloud_polling",
   "codeowners": [
     "@humbertogontijo"