diff --git a/homeassistant/helpers/update_coordinator.py b/homeassistant/helpers/update_coordinator.py index 16f3b9b696421..fcab7de352955 100644 --- a/homeassistant/helpers/update_coordinator.py +++ b/homeassistant/helpers/update_coordinator.py @@ -42,6 +42,16 @@ class UpdateFailed(HomeAssistantError): """Raised when an update has failed.""" + def __init__( + self, + *args: Any, + retry_after: float | None = None, + **kwargs: Any, + ) -> None: + """Initialize exception.""" + super().__init__(*args, **kwargs) + self.retry_after = retry_after + class BaseDataUpdateCoordinatorProtocol(Protocol): """Base protocol type for DataUpdateCoordinator.""" @@ -119,6 +129,7 @@ def __init__( self._unsub_refresh: CALLBACK_TYPE | None = None self._unsub_shutdown: CALLBACK_TYPE | None = None self._request_refresh_task: asyncio.TimerHandle | None = None + self._retry_after: float | None = None self.last_update_success = True self.last_exception: Exception | None = None @@ -250,9 +261,12 @@ def _schedule_refresh(self) -> None: hass = self.hass loop = hass.loop - next_refresh = ( - int(loop.time()) + self._microsecond + self._update_interval_seconds - ) + update_interval = self._update_interval_seconds + if self._retry_after is not None: + update_interval = self._retry_after + self._retry_after = None + + next_refresh = int(loop.time()) + self._microsecond + update_interval self._unsub_refresh = loop.call_at( next_refresh, self.__wrap_handle_refresh_interval ).cancel @@ -317,7 +331,9 @@ async def async_config_entry_first_refresh(self) -> None: ) if await self.__wrap_async_setup(): await self._async_refresh( - log_failures=False, raise_on_auth_failed=True, raise_on_entry_error=True + log_failures=False, + raise_on_auth_failed=True, + raise_on_entry_error=True, ) if self.last_update_success: return @@ -419,6 +435,16 @@ async def _async_refresh( # noqa: C901 except UpdateFailed as err: self.last_exception = err + # We can only honor a retry_after, after the config entry has been set up + # Basically meaning that the retry after can't be used when coming + # from an async_config_entry_first_refresh + if err.retry_after is not None and not raise_on_entry_error: + self._retry_after = err.retry_after + self.logger.debug( + "Retry after triggered. Scheduling next update in %s second(s)", + err.retry_after, + ) + if self.last_update_success: if log_failures: self.logger.error("Error fetching %s data: %s", self.name, err) diff --git a/tests/helpers/test_update_coordinator.py b/tests/helpers/test_update_coordinator.py index 57e80927e7ee6..069aec3544f2e 100644 --- a/tests/helpers/test_update_coordinator.py +++ b/tests/helpers/test_update_coordinator.py @@ -1053,3 +1053,106 @@ def stop_listen(self) -> None: # Ensure the coordinator is released assert weak_ref() is None + + +@pytest.mark.parametrize( + ("exc", "expected_exception", "message"), + [ + *KNOWN_ERRORS, + (Exception(), Exception, "Unknown exception"), + ( + update_coordinator.UpdateFailed(retry_after=60), + update_coordinator.UpdateFailed, + "Error fetching test data", + ), + ], +) +@pytest.mark.parametrize( + "method", + ["update_method", "setup_method"], +) +async def test_update_failed_retry_after( + hass: HomeAssistant, + exc: Exception, + expected_exception: type[Exception], + message: str, + method: str, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test async_config_entry_first_refresh raises ConfigEntryNotReady on failure. + + Verify we do not log the exception since raising ConfigEntryNotReady + will be caught by config_entries.async_setup which will log it with + a decreasing level of logging once the first message is logged. + """ + entry = MockConfigEntry() + entry.mock_state( + hass, + config_entries.ConfigEntryState.SETUP_IN_PROGRESS, + ) + crd = get_crd(hass, DEFAULT_UPDATE_INTERVAL, entry) + setattr(crd, method, AsyncMock(side_effect=exc)) + + with pytest.raises(ConfigEntryNotReady): + await crd.async_config_entry_first_refresh() + + assert crd.last_update_success is False + assert isinstance(crd.last_exception, expected_exception) + assert message not in caplog.text + + # Only to check the retry_after wasn't hit + assert crd._retry_after is None + + +@pytest.mark.parametrize( + ("exc", "expected_exception", "message"), + [ + ( + update_coordinator.UpdateFailed(retry_after=60), + update_coordinator.UpdateFailed, + "Error fetching test data", + ), + ], +) +async def test_refresh_known_errors_retry_after( + exc: update_coordinator.UpdateFailed, + expected_exception: type[Exception], + message: str, + crd: update_coordinator.DataUpdateCoordinator[int], + caplog: pytest.LogCaptureFixture, + hass: HomeAssistant, +) -> None: + """Test raising known errors, this time with retry_after.""" + unsub = crd.async_add_listener(lambda: None) + + crd.update_method = AsyncMock(side_effect=exc) + + with ( + patch.object(hass.loop, "time", return_value=1_000.0), + patch.object(hass.loop, "call_at") as mock_call_at, + ): + await crd.async_refresh() + + assert crd.data is None + assert crd.last_update_success is False + assert isinstance(crd.last_exception, expected_exception) + assert message in caplog.text + + when = mock_call_at.call_args[0][0] + + expected = 1_000.0 + crd._microsecond + exc.retry_after + assert abs(when - expected) < 0.005, (when, expected) + + assert crd._retry_after is None + + # Next schedule should fall back to regular update_interval + mock_call_at.reset_mock() + crd._schedule_refresh() + when2 = mock_call_at.call_args[0][0] + expected_cancelled = ( + 1_000.0 + crd._microsecond + crd.update_interval.total_seconds() + ) + assert abs(when2 - expected_cancelled) < 0.005, (when2, expected_cancelled) + + unsub() + crd._unschedule_refresh()