Add better error handling for Roborock initialization (#104181)

* Introduce better handling of errors in init for Roborock

* patch internally

* push exceptions up

* remove duplicated test

homeassistant/components/roborock/__init__.py

@@ -2,13 +2,15 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Coroutine
 from datetime import timedelta
 import logging
+from typing import Any
 
 from roborock import RoborockException, RoborockInvalidCredentials
 from roborock.api import RoborockApiClient
 from roborock.cloud_api import RoborockMqttClient
-from roborock.containers import DeviceData, HomeDataDevice, UserData
+from roborock.containers import DeviceData, HomeDataDevice, HomeDataProduct, UserData
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_USERNAME
@@ -40,61 +42,103 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     device_map: dict[str, HomeDataDevice] = {
         device.duid: device for device in home_data.devices + home_data.received_devices
     }
-    product_info = {product.id: product for product in home_data.products}
+    product_info: dict[str, HomeDataProduct] = {
-    # Create a mqtt_client, which is needed to get the networking information of the device for local connection and in the future, get the map.
+        product.id: product for product in home_data.products
-    mqtt_clients = {
-        device.duid: RoborockMqttClient(
-            user_data, DeviceData(device, product_info[device.product_id].model)
-        )
-        for device in device_map.values()
     }
-    network_results = await asyncio.gather(
+    # Get a Coordinator if the device is available or if we have connected to the device before
-        *(mqtt_client.get_networking() for mqtt_client in mqtt_clients.values())
+    coordinators = await asyncio.gather(
-    )
+        *build_setup_functions(hass, device_map, user_data, product_info),
-    network_info = {
-        device.duid: result
-        for device, result in zip(device_map.values(), network_results)
-        if result is not None
-    }
-    if not network_info:
-        raise ConfigEntryNotReady(
-            "Could not get network information about your devices"
-        )
-    coordinator_map: dict[str, RoborockDataUpdateCoordinator] = {}
-    for device_id, device in device_map.items():
-        coordinator_map[device_id] = RoborockDataUpdateCoordinator(
-            hass,
-            device,
-            network_info[device_id],
-            product_info[device.product_id],
-            mqtt_clients[device.duid],
-        )
-    await asyncio.gather(
-        *(coordinator.verify_api() for coordinator in coordinator_map.values())
-    )
-    # If one device update fails - we still want to set up other devices
-    await asyncio.gather(
-        *(
-            coordinator.async_config_entry_first_refresh()
-            for coordinator in coordinator_map.values()
-        ),
         return_exceptions=True,
     )
+    # Valid coordinators are those where we had networking cached or we could get networking
+    valid_coordinators: list[RoborockDataUpdateCoordinator] = [
+        coord
+        for coord in coordinators
+        if isinstance(coord, RoborockDataUpdateCoordinator)
+    ]
+    if len(valid_coordinators) == 0:
+        raise ConfigEntryNotReady("No coordinators were able to successfully setup.")
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {
-        device_id: coordinator
+        coordinator.roborock_device_info.device.duid: coordinator
-        for device_id, coordinator in coordinator_map.items()
+        for coordinator in valid_coordinators
-        if coordinator.last_update_success
+    }
-    }  # Only add coordinators that succeeded
-
-    if not hass.data[DOMAIN][entry.entry_id]:
-        # Don't start if no coordinators succeeded.
-        raise ConfigEntryNotReady("There are no devices that can currently be reached.")
-
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
 
 
+def build_setup_functions(
+    hass: HomeAssistant,
+    device_map: dict[str, HomeDataDevice],
+    user_data: UserData,
+    product_info: dict[str, HomeDataProduct],
+) -> list[Coroutine[Any, Any, RoborockDataUpdateCoordinator | None]]:
+    """Create a list of setup functions that can later be called asynchronously."""
+    setup_functions = []
+    for device in device_map.values():
+        setup_functions.append(
+            setup_device(hass, user_data, device, product_info[device.product_id])
+        )
+    return setup_functions
+
+
+async def setup_device(
+    hass: HomeAssistant,
+    user_data: UserData,
+    device: HomeDataDevice,
+    product_info: HomeDataProduct,
+) -> RoborockDataUpdateCoordinator | None:
+    """Set up a device Coordinator."""
+    mqtt_client = RoborockMqttClient(user_data, DeviceData(device, product_info.name))
+    try:
+        networking = await mqtt_client.get_networking()
+        if networking is None:
+            # If the api does not return an error but does return None for
+            # get_networking - then we need to go through cache checking.
+            raise RoborockException("Networking request returned None.")
+    except RoborockException as err:
+        _LOGGER.warning(
+            "Not setting up %s because we could not get the network information of the device. "
+            "Please confirm it is online and the Roborock servers can communicate with it",
+            device.name,
+        )
+        _LOGGER.debug(err)
+        raise err
+    coordinator = RoborockDataUpdateCoordinator(
+        hass, device, networking, product_info, mqtt_client
+    )
+    # Verify we can communicate locally - if we can't, switch to cloud api
+    await coordinator.verify_api()
+    coordinator.api.is_available = True
+    try:
+        await coordinator.async_config_entry_first_refresh()
+    except ConfigEntryNotReady:
+        if isinstance(coordinator.api, RoborockMqttClient):
+            _LOGGER.warning(
+                "Not setting up %s because the we failed to get data for the first time using the online client. "
+                "Please ensure your Home Assistant instance can communicate with this device. "
+                "You may need to open firewall instances on your Home Assistant network and on your Vacuum's network",
+                device.name,
+            )
+            # Most of the time if we fail to connect using the mqtt client, the problem is due to firewall,
+            # but in case if it isn't, the error can be included in debug logs for the user to grab.
+            if coordinator.last_exception:
+                _LOGGER.debug(coordinator.last_exception)
+                raise coordinator.last_exception
+        elif coordinator.last_exception:
+            # If this is reached, we have verified that we can communicate with the Vacuum locally,
+            # so if there is an error here - it is not a communication issue but some other problem
+            extra_error = f"Please create an issue with the following error included: {coordinator.last_exception}"
+            _LOGGER.warning(
+                "Not setting up %s because the coordinator failed to get data for the first time using the "
+                "offline client %s",
+                device.name,
+                extra_error,
+            )
+            raise coordinator.last_exception
+    return coordinator
+
+
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Handle removal of an entry."""
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

tests/components/roborock/test_init.py

@@ -6,7 +6,6 @@ from roborock import RoborockException, RoborockInvalidCredentials
 from homeassistant.components.roborock.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.update_coordinator import UpdateFailed
 from homeassistant.setup import async_setup_component
 
 from tests.common import MockConfigEntry
@@ -35,8 +34,74 @@ async def test_config_entry_not_ready(
     with patch(
         "homeassistant.components.roborock.RoborockApiClient.get_home_data",
     ), patch(
-        "homeassistant.components.roborock.RoborockDataUpdateCoordinator._async_update_data",
+        "homeassistant.components.roborock.coordinator.RoborockLocalClient.get_prop",
-        side_effect=UpdateFailed(),
+        side_effect=RoborockException(),
+    ):
+        await async_setup_component(hass, DOMAIN, {})
+        assert mock_roborock_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_config_entry_not_ready_home_data(
+    hass: HomeAssistant, mock_roborock_entry: MockConfigEntry
+) -> None:
+    """Test that when we fail to get home data, entry retries."""
+    with patch(
+        "homeassistant.components.roborock.RoborockApiClient.get_home_data",
+        side_effect=RoborockException(),
+    ), patch(
+        "homeassistant.components.roborock.coordinator.RoborockLocalClient.get_prop",
+        side_effect=RoborockException(),
+    ):
+        await async_setup_component(hass, DOMAIN, {})
+        assert mock_roborock_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_get_networking_fails(
+    hass: HomeAssistant, mock_roborock_entry: MockConfigEntry, bypass_api_fixture
+) -> None:
+    """Test that when networking fails, we attempt to retry."""
+    with patch(
+        "homeassistant.components.roborock.RoborockMqttClient.get_networking",
+        side_effect=RoborockException(),
+    ):
+        await async_setup_component(hass, DOMAIN, {})
+        assert mock_roborock_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_get_networking_fails_none(
+    hass: HomeAssistant, mock_roborock_entry: MockConfigEntry, bypass_api_fixture
+) -> None:
+    """Test that when networking returns None, we attempt to retry."""
+    with patch(
+        "homeassistant.components.roborock.RoborockMqttClient.get_networking",
+        return_value=None,
+    ):
+        await async_setup_component(hass, DOMAIN, {})
+        assert mock_roborock_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_cloud_client_fails_props(
+    hass: HomeAssistant, mock_roborock_entry: MockConfigEntry, bypass_api_fixture
+) -> None:
+    """Test that if networking succeeds, but we can't communicate with the vacuum, we can't get props, fail."""
+    with patch(
+        "homeassistant.components.roborock.coordinator.RoborockLocalClient.ping",
+        side_effect=RoborockException(),
+    ), patch(
+        "homeassistant.components.roborock.coordinator.RoborockMqttClient.get_prop",
+        side_effect=RoborockException(),
+    ):
+        await async_setup_component(hass, DOMAIN, {})
+        assert mock_roborock_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_local_client_fails_props(
+    hass: HomeAssistant, mock_roborock_entry: MockConfigEntry, bypass_api_fixture
+) -> None:
+    """Test that if networking succeeds, but we can't communicate locally with the vacuum, we can't get props, fail."""
+    with patch(
+        "homeassistant.components.roborock.coordinator.RoborockLocalClient.get_prop",
+        side_effect=RoborockException(),
     ):
         await async_setup_component(hass, DOMAIN, {})
         assert mock_roborock_entry.state is ConfigEntryState.SETUP_RETRY
@@ -55,15 +120,3 @@ async def test_reauth_started(
     flows = hass.config_entries.flow.async_progress()
     assert len(flows) == 1
     assert flows[0]["step_id"] == "reauth_confirm"
-
-
-async def test_config_entry_not_ready_home_data(
-    hass: HomeAssistant, mock_roborock_entry: MockConfigEntry
-) -> None:
-    """Test that when we fail to get home data, entry retries."""
-    with patch(
-        "homeassistant.components.roborock.RoborockApiClient.get_home_data",
-        side_effect=RoborockException(),
-    ):
-        await async_setup_component(hass, DOMAIN, {})
-        assert mock_roborock_entry.state is ConfigEntryState.SETUP_RETRY