From fb04df5392cdc80266b9b9998dcfb46d0ec1d0ac Mon Sep 17 00:00:00 2001 From: Christopher Bailey Date: Wed, 21 Feb 2024 00:31:42 -0500 Subject: [PATCH] Improve UniFi Protect re-auth (#110021) Co-authored-by: J. Nick Koston --- .../components/unifiprotect/__init__.py | 7 ++++ .../components/unifiprotect/const.py | 3 ++ homeassistant/components/unifiprotect/data.py | 3 +- tests/components/unifiprotect/test_init.py | 42 +++++++++++-------- 4 files changed, 37 insertions(+), 18 deletions(-) diff --git a/homeassistant/components/unifiprotect/__init__.py b/homeassistant/components/unifiprotect/__init__.py index 8639b0becdc..4d060161fae 100644 --- a/homeassistant/components/unifiprotect/__init__.py +++ b/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 diff --git a/homeassistant/components/unifiprotect/const.py b/homeassistant/components/unifiprotect/const.py index 3bc689666c7..b3bba4ece39 100644 --- a/homeassistant/components/unifiprotect/const.py +++ b/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" diff --git a/homeassistant/components/unifiprotect/data.py b/homeassistant/components/unifiprotect/data.py index 11782c42bee..5565581290c 100644 --- a/homeassistant/components/unifiprotect/data.py +++ b/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: diff --git a/tests/components/unifiprotect/test_init.py b/tests/components/unifiprotect/test_init.py index caa77e8408d..d1dd4379515 100644 --- a/tests/components/unifiprotect/test_init.py +++ b/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.entry.state == ConfigEntryState.LOADED + assert ufp.api.update.call_count == AUTH_RETRIES + 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