Add ability to shutdown update coordinator (#91456)

* Add ability to shutdown update coordinator

* Adjust nibe_heatpump

* Add tests

* Use async

* Remove duplicate code in update coordinator

* Adjust

* Revert nibe changes - it can now be done in a follow-up PR

* Adjust

* Fix incorrect merge

* async_fire_time_changed

homeassistant/helpers/update_coordinator.py

@@ -71,6 +71,7 @@ class DataUpdateCoordinator(BaseDataUpdateCoordinatorProtocol, Generic[_T]):
         self.name = name
         self.update_method = update_method
         self.update_interval = update_interval
+        self._shutdown_requested = False
         self.config_entry = config_entries.current_entry.get()
 
         # It's None before the first successful update.
@@ -141,6 +142,12 @@ class DataUpdateCoordinator(BaseDataUpdateCoordinatorProtocol, Generic[_T]):
         for update_callback, _ in list(self._listeners.values()):
             update_callback()
 
+    async def async_shutdown(self) -> None:
+        """Cancel any scheduled call, and ignore new runs."""
+        self._shutdown_requested = True
+        self._async_unsub_refresh()
+        await self._debounced_refresh.async_shutdown()
+
     @callback
     def _unschedule_refresh(self) -> None:
         """Unschedule any pending refresh since there is no longer any listeners."""
@@ -237,7 +244,7 @@ class DataUpdateCoordinator(BaseDataUpdateCoordinatorProtocol, Generic[_T]):
         self._async_unsub_refresh()
         self._debounced_refresh.async_cancel()
 
-        if scheduled and self.hass.is_stopping:
+        if self._shutdown_requested or scheduled and self.hass.is_stopping:
             return
 
         if log_timing := self.logger.isEnabledFor(logging.DEBUG):

tests/helpers/test_update_coordinator.py

@@ -116,6 +116,43 @@ async def test_async_refresh(
     assert updates == [2]
 
 
+async def test_shutdown(
+    hass: HomeAssistant,
+    crd: update_coordinator.DataUpdateCoordinator[int],
+) -> None:
+    """Test async_shutdown for update coordinator."""
+    assert crd.data is None
+    await crd.async_refresh()
+    assert crd.data == 1
+    assert crd.last_update_success is True
+    # Make sure we didn't schedule a refresh because we have 0 listeners
+    assert crd._unsub_refresh is None
+
+    updates = []
+
+    def update_callback():
+        updates.append(crd.data)
+
+    _ = crd.async_add_listener(update_callback)
+    await crd.async_refresh()
+    assert updates == [2]
+    assert crd._unsub_refresh is not None
+
+    # Test shutdown through function
+    with patch.object(crd._debounced_refresh, "async_shutdown") as mock_shutdown:
+        await crd.async_shutdown()
+
+    async_fire_time_changed(hass, utcnow() + crd.update_interval)
+    await hass.async_block_till_done()
+
+    # Test we shutdown the debouncer and cleared the subscriptions
+    assert len(mock_shutdown.mock_calls) == 1
+    assert crd._unsub_refresh is None
+
+    await crd.async_refresh()
+    assert updates == [2]
+
+
 async def test_update_context(
     crd: update_coordinator.DataUpdateCoordinator[int],
 ) -> None: