Improve setup exception handling in ruckus_unleashed (#168014)

homeassistant/components/ruckus_unleashed/__init__.py

@@ -49,10 +49,14 @@ async def async_setup_entry(
 
     await coordinator.async_config_entry_first_refresh()
 
-    system_info = await ruckus.api.get_system_info()
+    try:
+        system_info = await ruckus.api.get_system_info()
+        aps = await ruckus.api.get_aps()
+    except (ConnectionError, SchemaError) as err:
+        await ruckus.close()
+        raise ConfigEntryNotReady from err
 
     registry = dr.async_get(hass)
-    aps = await ruckus.api.get_aps()
     for access_point in aps:
         _LOGGER.debug("AP [%s] %s", access_point[API_AP_MAC], entry.entry_id)
         registry.async_get_or_create(

tests/components/ruckus_unleashed/test_init.py

@@ -1,9 +1,11 @@
 """Test the Ruckus config flow."""
 
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, patch
 
+from aioruckus import RuckusAjaxApi
 from aioruckus.const import ERROR_CONNECT_TIMEOUT, ERROR_LOGIN_INCORRECT
-from aioruckus.exceptions import AuthenticationError
+from aioruckus.exceptions import AuthenticationError, SchemaError
+import pytest
 
 from homeassistant.components.ruckus_unleashed.const import (
     API_AP_DEVNAME,
@@ -90,3 +92,30 @@ async def test_unload_entry(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
 
     assert entry.state is ConfigEntryState.NOT_LOADED
+
+
+@pytest.mark.parametrize(
+    ("method", "error"),
+    [
+        ("get_system_info", ConnectionError("connection lost")),
+        ("get_aps", SchemaError("unexpected schema")),
+    ],
+)
+async def test_setup_entry_error_post_login(
+    hass: HomeAssistant, method: str, error: Exception
+) -> None:
+    """Test entry setup retries on post-login API errors."""
+    entry = mock_config_entry()
+    entry.add_to_hass(hass)
+    with (
+        RuckusAjaxApiPatchContext(),
+        patch.object(
+            RuckusAjaxApi,
+            method,
+            new=AsyncMock(side_effect=error),
+        ),
+    ):
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.SETUP_RETRY