Add retry_after to UpdateFailed in update coordinator

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)

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()