Improve UniFi Protect re-auth (#110021)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/unifiprotect/__init__.py

@@ -20,6 +20,7 @@ from homeassistant.helpers.issue_registry import IssueSeverity
 from homeassistant.helpers.typing import ConfigType
 
 from .const import (
+    AUTH_RETRIES,
     CONF_ALLOW_EA,
     DEFAULT_SCAN_INTERVAL,
     DEVICES_THAT_ADOPT,
@@ -62,6 +63,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     try:
         nvr_info = await protect.get_nvr()
     except NotAuthorized as err:
+        retry_key = f"{entry.entry_id}_auth"
+        retries = hass.data.setdefault(DOMAIN, {}).get(retry_key, 0)
+        if retries < AUTH_RETRIES:
+            retries += 1
+            hass.data[DOMAIN][retry_key] = retries
+            raise ConfigEntryNotReady from err
         raise ConfigEntryAuthFailed(err) from err
     except (TimeoutError, ClientError, ServerDisconnectedError) as err:
         raise ConfigEntryNotReady from err

homeassistant/components/unifiprotect/const.py

@@ -5,6 +5,9 @@ from pyunifiprotect.data import ModelType, Version
 from homeassistant.const import Platform
 
 DOMAIN = "unifiprotect"
+# some UniFi OS consoles have an unknown rate limit on auth
+# if rate limit is triggered a 401 is returned
+AUTH_RETRIES = 11  # ~12 hours of retries with the last waiting ~6 hours
 
 ATTR_EVENT_SCORE = "event_score"
 ATTR_EVENT_ID = "event_id"

homeassistant/components/unifiprotect/data.py

@@ -27,6 +27,7 @@ from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.event import async_track_time_interval
 
 from .const import (
+    AUTH_RETRIES,
     CONF_DISABLE_RTSP,
     CONF_MAX_MEDIA,
     DEFAULT_MAX_MEDIA,
@@ -133,7 +134,7 @@ class ProtectData:
         try:
             updates = await self.api.update(force=force)
         except NotAuthorized:
-            if self._auth_failures < 10:
+            if self._auth_failures < AUTH_RETRIES:
                 _LOGGER.exception("Auth error while updating")
                 self._auth_failures += 1
             else:

tests/components/unifiprotect/test_init.py

@@ -9,6 +9,7 @@ from pyunifiprotect import NotAuthorized, NvrError, ProtectApiClient
 from pyunifiprotect.data import NVR, Bootstrap, Light
 
 from homeassistant.components.unifiprotect.const import (
+    AUTH_RETRIES,
     CONF_DISABLE_RTSP,
     DEFAULT_SCAN_INTERVAL,
     DOMAIN,
@@ -47,7 +48,7 @@ async def test_setup(hass: HomeAssistant, ufp: MockUFPFixture) -> None:
     await hass.config_entries.async_setup(ufp.entry.entry_id)
     await hass.async_block_till_done()
 
-    assert ufp.entry.state == ConfigEntryState.LOADED
+    assert ufp.entry.state is ConfigEntryState.LOADED
     assert ufp.api.update.called
     assert ufp.entry.unique_id == ufp.api.bootstrap.nvr.mac
 
@@ -62,7 +63,7 @@ async def test_setup_multiple(
     await hass.config_entries.async_setup(ufp.entry.entry_id)
     await hass.async_block_till_done()
 
-    assert ufp.entry.state == ConfigEntryState.LOADED
+    assert ufp.entry.state is ConfigEntryState.LOADED
     assert ufp.api.update.called
     assert ufp.entry.unique_id == ufp.api.bootstrap.nvr.mac
 
@@ -104,14 +105,14 @@ async def test_reload(hass: HomeAssistant, ufp: MockUFPFixture) -> None:
 
     await hass.config_entries.async_setup(ufp.entry.entry_id)
     await hass.async_block_till_done()
-    assert ufp.entry.state == ConfigEntryState.LOADED
+    assert ufp.entry.state is ConfigEntryState.LOADED
 
     options = dict(ufp.entry.options)
     options[CONF_DISABLE_RTSP] = True
     hass.config_entries.async_update_entry(ufp.entry, options=options)
     await hass.async_block_till_done()
 
-    assert ufp.entry.state == ConfigEntryState.LOADED
+    assert ufp.entry.state is ConfigEntryState.LOADED
     assert ufp.api.async_disconnect_ws.called
 
 
@@ -119,10 +120,10 @@ async def test_unload(hass: HomeAssistant, ufp: MockUFPFixture, light: Light) ->
     """Test unloading of unifiprotect entry."""
 
     await init_entry(hass, ufp, [light])
-    assert ufp.entry.state == ConfigEntryState.LOADED
+    assert ufp.entry.state is ConfigEntryState.LOADED
 
     await hass.config_entries.async_unload(ufp.entry.entry_id)
-    assert ufp.entry.state == ConfigEntryState.NOT_LOADED
+    assert ufp.entry.state is ConfigEntryState.NOT_LOADED
     assert ufp.api.async_disconnect_ws.called
 
 
@@ -135,7 +136,7 @@ async def test_setup_too_old(
 
     await hass.config_entries.async_setup(ufp.entry.entry_id)
     await hass.async_block_till_done()
-    assert ufp.entry.state == ConfigEntryState.SETUP_ERROR
+    assert ufp.entry.state is ConfigEntryState.SETUP_ERROR
     assert not ufp.api.update.called
 
 
@@ -146,7 +147,7 @@ async def test_setup_failed_update(hass: HomeAssistant, ufp: MockUFPFixture) ->
 
     await hass.config_entries.async_setup(ufp.entry.entry_id)
     await hass.async_block_till_done()
-    assert ufp.entry.state == ConfigEntryState.SETUP_RETRY
+    assert ufp.entry.state is ConfigEntryState.SETUP_RETRY
     assert ufp.api.update.called
 
 
@@ -157,20 +158,20 @@ async def test_setup_failed_update_reauth(
 
     await hass.config_entries.async_setup(ufp.entry.entry_id)
     await hass.async_block_till_done()
-    assert ufp.entry.state == ConfigEntryState.LOADED
+    assert ufp.entry.state is ConfigEntryState.LOADED
 
     # reauth should not be triggered until there are 10 auth failures in a row
     # to verify it is not transient
     ufp.api.update = AsyncMock(side_effect=NotAuthorized)
-    for _ in range(10):
+    for _ in range(AUTH_RETRIES):
         await time_changed(hass, DEFAULT_SCAN_INTERVAL)
         assert len(hass.config_entries.flow._progress) == 0
 
-    assert ufp.api.update.call_count == 10
+    assert ufp.api.update.call_count == AUTH_RETRIES
-    assert ufp.entry.state == ConfigEntryState.LOADED
+    assert ufp.entry.state is ConfigEntryState.LOADED
 
     await time_changed(hass, DEFAULT_SCAN_INTERVAL)
-    assert ufp.api.update.call_count == 11
+    assert ufp.api.update.call_count == AUTH_RETRIES + 1
     assert len(hass.config_entries.flow._progress) == 1
 
 
@@ -181,17 +182,24 @@ async def test_setup_failed_error(hass: HomeAssistant, ufp: MockUFPFixture) -> N
 
     await hass.config_entries.async_setup(ufp.entry.entry_id)
     await hass.async_block_till_done()
-    assert ufp.entry.state == ConfigEntryState.SETUP_RETRY
+    assert ufp.entry.state is ConfigEntryState.SETUP_RETRY
     assert not ufp.api.update.called
 
 
 async def test_setup_failed_auth(hass: HomeAssistant, ufp: MockUFPFixture) -> None:
-    """Test setup of unifiprotect entry with unauthorized error."""
+    """Test setup of unifiprotect entry with unauthorized error after multiple retries."""
 
     ufp.api.get_nvr = AsyncMock(side_effect=NotAuthorized)
 
     await hass.config_entries.async_setup(ufp.entry.entry_id)
-    assert ufp.entry.state == ConfigEntryState.SETUP_ERROR
+    assert ufp.entry.state is ConfigEntryState.SETUP_RETRY
+
+    for _ in range(AUTH_RETRIES - 1):
+        await hass.config_entries.async_reload(ufp.entry.entry_id)
+        assert ufp.entry.state is ConfigEntryState.SETUP_RETRY
+
+    await hass.config_entries.async_reload(ufp.entry.entry_id)
+    assert ufp.entry.state is ConfigEntryState.SETUP_ERROR
     assert not ufp.api.update.called
 
 
@@ -208,7 +216,7 @@ async def test_setup_starts_discovery(
 
         await hass.config_entries.async_setup(ufp.entry.entry_id)
         await hass.async_block_till_done()
-        assert ufp.entry.state == ConfigEntryState.LOADED
+        assert ufp.entry.state is ConfigEntryState.LOADED
         await hass.async_block_till_done()
         assert len(hass.config_entries.flow.async_progress_by_handler(DOMAIN)) == 1