diff --git a/homeassistant/helpers/update_coordinator.py b/homeassistant/helpers/update_coordinator.py index c563ef09a52..5706e34dc9c 100644 --- a/homeassistant/helpers/update_coordinator.py +++ b/homeassistant/helpers/update_coordinator.py @@ -15,7 +15,8 @@ import aiohttp 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 @@ class DataUpdateCoordinator(BaseDataUpdateCoordinatorProtocol, Generic[_T]): 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 @@ class DataUpdateCoordinator(BaseDataUpdateCoordinatorProtocol, Generic[_T]): 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 @@ class DataUpdateCoordinator(BaseDataUpdateCoordinatorProtocol, Generic[_T]): """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 @@ class DataUpdateCoordinator(BaseDataUpdateCoordinatorProtocol, Generic[_T]): 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.""" diff --git a/tests/helpers/test_update_coordinator.py b/tests/helpers/test_update_coordinator.py index 35ae78834c5..91f761b5bb6 100644 --- a/tests/helpers/test_update_coordinator.py +++ b/tests/helpers/test_update_coordinator.py @@ -186,6 +186,38 @@ async def test_shutdown_on_entry_unload( 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: