From f006716173eb16848b0bcf53d8309fcdea803019 Mon Sep 17 00:00:00 2001 From: Josef Zweck <24647999+zweckj@users.noreply.github.com> Date: Fri, 19 Jul 2024 14:24:25 +0200 Subject: [PATCH] Add `async_setup` method to `DataUpdateCoordinator` (#116677) * init * Update homeassistant/helpers/update_coordinator.py Co-authored-by: Joost Lekkerkerker * fix typo, ruff * consistency with rest, test * pylint suppression * ruff * ruff * switch to one test * add last exc * add tests for auth & Entry Errors * move exceptions to correct test * Update update_coordinator.py Co-authored-by: G Johansson * test setup call * simplify --------- Co-authored-by: Joost Lekkerkerker Co-authored-by: G Johansson --- homeassistant/helpers/update_coordinator.py | 53 +++++++++++++++++--- tests/helpers/test_update_coordinator.py | 54 +++++++++++++++++++-- 2 files changed, 98 insertions(+), 9 deletions(-) diff --git a/homeassistant/helpers/update_coordinator.py b/homeassistant/helpers/update_coordinator.py index 7cb1082d34d..4fe4953d752 100644 --- a/homeassistant/helpers/update_coordinator.py +++ b/homeassistant/helpers/update_coordinator.py @@ -71,6 +71,7 @@ class DataUpdateCoordinator(BaseDataUpdateCoordinatorProtocol, Generic[_DataT]): name: str, update_interval: timedelta | None = None, update_method: Callable[[], Awaitable[_DataT]] | None = None, + setup_method: Callable[[], Awaitable[None]] | None = None, request_refresh_debouncer: Debouncer[Coroutine[Any, Any, None]] | None = None, always_update: bool = True, ) -> None: @@ -79,6 +80,7 @@ class DataUpdateCoordinator(BaseDataUpdateCoordinatorProtocol, Generic[_DataT]): self.logger = logger self.name = name self.update_method = update_method + self.setup_method = setup_method self._update_interval_seconds: float | None = None self.update_interval = update_interval self._shutdown_requested = False @@ -275,15 +277,54 @@ class DataUpdateCoordinator(BaseDataUpdateCoordinatorProtocol, Generic[_DataT]): fails. Additionally logging is handled by config entry setup to ensure that multiple retries do not cause log spam. """ - await self._async_refresh( - log_failures=False, raise_on_auth_failed=True, raise_on_entry_error=True - ) - if self.last_update_success: - return + if await self.__wrap_async_setup(): + await self._async_refresh( + log_failures=False, raise_on_auth_failed=True, raise_on_entry_error=True + ) + if self.last_update_success: + return ex = ConfigEntryNotReady() ex.__cause__ = self.last_exception raise ex + async def __wrap_async_setup(self) -> bool: + """Error handling for _async_setup.""" + try: + await self._async_setup() + except ( + TimeoutError, + requests.exceptions.Timeout, + aiohttp.ClientError, + requests.exceptions.RequestException, + urllib.error.URLError, + UpdateFailed, + ) as err: + self.last_exception = err + + except (ConfigEntryError, ConfigEntryAuthFailed) as err: + self.last_exception = err + self.last_update_success = False + raise + + except Exception as err: # pylint: disable=broad-except + self.last_exception = err + self.logger.exception("Unexpected error fetching %s data", self.name) + else: + return True + + self.last_update_success = False + return False + + async def _async_setup(self) -> None: + """Set up the coordinator. + + Can be overwritten by integrations to load data or resources + only once during the first refresh. + """ + if self.setup_method is None: + return None + return await self.setup_method() + async def async_refresh(self) -> None: """Refresh data and log errors.""" await self._async_refresh(log_failures=True) @@ -393,7 +434,7 @@ class DataUpdateCoordinator(BaseDataUpdateCoordinatorProtocol, Generic[_DataT]): self.logger.debug( "Finished fetching %s data in %.3f seconds (success: %s)", self.name, - monotonic() - start, + monotonic() - start, # pylint: disable=possibly-used-before-assignment self.last_update_success, ) if not auth_failed and self._listeners and not self.hass.is_stopping: diff --git a/tests/helpers/test_update_coordinator.py b/tests/helpers/test_update_coordinator.py index 8633bf862a5..d450d924f1f 100644 --- a/tests/helpers/test_update_coordinator.py +++ b/tests/helpers/test_update_coordinator.py @@ -13,7 +13,11 @@ import requests from homeassistant import config_entries from homeassistant.const import EVENT_HOMEASSISTANT_STOP from homeassistant.core import CoreState, HomeAssistant, callback -from homeassistant.exceptions import ConfigEntryNotReady +from homeassistant.exceptions import ( + ConfigEntryAuthFailed, + ConfigEntryError, + ConfigEntryNotReady, +) from homeassistant.helpers import update_coordinator from homeassistant.util.dt import utcnow @@ -525,11 +529,19 @@ async def test_stop_refresh_on_ha_stop( @pytest.mark.parametrize( "err_msg", - KNOWN_ERRORS, + [ + *KNOWN_ERRORS, + (Exception(), Exception, "Unknown exception"), + ], +) +@pytest.mark.parametrize( + "method", + ["update_method", "setup_method"], ) async def test_async_config_entry_first_refresh_failure( err_msg: tuple[Exception, type[Exception], str], crd: update_coordinator.DataUpdateCoordinator[int], + method: str, caplog: pytest.LogCaptureFixture, ) -> None: """Test async_config_entry_first_refresh raises ConfigEntryNotReady on failure. @@ -538,7 +550,7 @@ async def test_async_config_entry_first_refresh_failure( will be caught by config_entries.async_setup which will log it with a decreasing level of logging once the first message is logged. """ - crd.update_method = AsyncMock(side_effect=err_msg[0]) + setattr(crd, method, AsyncMock(side_effect=err_msg[0])) with pytest.raises(ConfigEntryNotReady): await crd.async_config_entry_first_refresh() @@ -548,13 +560,49 @@ async def test_async_config_entry_first_refresh_failure( assert err_msg[2] not in caplog.text +@pytest.mark.parametrize( + "err_msg", + [ + (ConfigEntryError(), ConfigEntryError, "Config entry error"), + (ConfigEntryAuthFailed(), ConfigEntryAuthFailed, "Config entry error"), + ], +) +@pytest.mark.parametrize( + "method", + ["update_method", "setup_method"], +) +async def test_async_config_entry_first_refresh_failure_passed_through( + err_msg: tuple[Exception, type[Exception], str], + crd: update_coordinator.DataUpdateCoordinator[int], + method: str, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test async_config_entry_first_refresh passes through ConfigEntryError & ConfigEntryAuthFailed. + + Verify we do not log the exception since it + will be caught by config_entries.async_setup which will log it with + a decreasing level of logging once the first message is logged. + """ + setattr(crd, method, AsyncMock(side_effect=err_msg[0])) + + with pytest.raises(err_msg[1]): + await crd.async_config_entry_first_refresh() + + assert crd.last_update_success is False + assert isinstance(crd.last_exception, err_msg[1]) + assert err_msg[2] not in caplog.text + + async def test_async_config_entry_first_refresh_success( crd: update_coordinator.DataUpdateCoordinator[int], caplog: pytest.LogCaptureFixture ) -> None: """Test first refresh successfully.""" + + crd.setup_method = AsyncMock() await crd.async_config_entry_first_refresh() assert crd.last_update_success is True + crd.setup_method.assert_called_once() async def test_not_schedule_refresh_if_system_option_disable_polling(