Add ability to shutdown a coordinator on STOP (#92611)

homeassistant/helpers/update_coordinator.py

@@ -15,7 +15,8 @@
 import requests
 
 from homeassistant import config_entries
-from homeassistant.core import CALLBACK_TYPE, HassJob, HomeAssistant, callback
+from homeassistant.const import EVENT_HOMEASSISTANT_STOP
+from homeassistant.core import CALLBACK_TYPE, Event, HassJob, HomeAssistant, callback
 from homeassistant.exceptions import (
     ConfigEntryAuthFailed,
     ConfigEntryError,
@@ -97,6 +98,7 @@ def __init__(
             job_name += f" {entry.title} {entry.domain} {entry.entry_id}"
         self._job = HassJob(self._handle_refresh_interval, job_name)
         self._unsub_refresh: CALLBACK_TYPE | None = None
+        self._unsub_shutdown: CALLBACK_TYPE | None = None
         self._request_refresh_task: asyncio.TimerHandle | None = None
         self.last_update_success = True
         self.last_exception: Exception | None = None
@@ -117,6 +119,22 @@ def __init__(
         if self.config_entry:
             self.config_entry.async_on_unload(self.async_shutdown)
 
+    async def async_register_shutdown(self) -> None:
+        """Register shutdown on HomeAssistant stop.
+
+        Should only be used by coordinators that are not linked to a config entry.
+        """
+        if self.config_entry:
+            raise RuntimeError("This should only be used outside of config entries.")
+
+        async def _on_hass_stop(_: Event) -> None:
+            """Shutdown coordinator on HomeAssistant stop."""
+            await self.async_shutdown()
+
+        self._unsub_shutdown = self.hass.bus.async_listen_once(
+            EVENT_HOMEASSISTANT_STOP, _on_hass_stop
+        )
+
     @callback
     def async_add_listener(
         self, update_callback: CALLBACK_TYPE, context: Any = None
@@ -149,6 +167,7 @@ async def async_shutdown(self) -> None:
         """Cancel any scheduled call, and ignore new runs."""
         self._shutdown_requested = True
         self._async_unsub_refresh()
+        self._async_unsub_shutdown()
         await self._debounced_refresh.async_shutdown()
 
     @callback
@@ -169,6 +188,12 @@ def _async_unsub_refresh(self) -> None:
             self._unsub_refresh()
             self._unsub_refresh = None
 
+    def _async_unsub_shutdown(self) -> None:
+        """Cancel any scheduled call."""
+        if self._unsub_shutdown:
+            self._unsub_shutdown()
+            self._unsub_shutdown = None
+
     @callback
     def _schedule_refresh(self) -> None:
         """Schedule a refresh."""

tests/helpers/test_update_coordinator.py

@@ -186,6 +186,38 @@ async def _refresh() -> int:
     assert crd._unsub_refresh is None
 
 
+async def test_shutdown_on_hass_stop(
+    hass: HomeAssistant,
+    crd: update_coordinator.DataUpdateCoordinator[int],
+) -> None:
+    """Test shutdown can be shutdown on STOP event."""
+    calls = 0
+
+    async def _refresh() -> int:
+        nonlocal calls
+        calls += 1
+        return calls
+
+    crd = update_coordinator.DataUpdateCoordinator[int](
+        hass,
+        _LOGGER,
+        name="test",
+        update_method=_refresh,
+        update_interval=DEFAULT_UPDATE_INTERVAL,
+    )
+    await crd.async_register_shutdown()
+
+    crd.async_add_listener(lambda: None)
+    assert crd._unsub_refresh is not None
+    assert not crd._shutdown_requested
+
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
+    await hass.async_block_till_done()
+
+    assert crd._shutdown_requested
+    assert crd._unsub_refresh is None
+
+
 async def test_update_context(
     crd: update_coordinator.DataUpdateCoordinator[int],
 ) -> None: