Skip to content

Commit

Permalink
Add ability to shutdown update coordinator (home-assistant#91456)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
epenet authored Apr 18, 2023
1 parent bdffb1f commit ae0cbff
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 1 deletion.
9 changes: 8 additions & 1 deletion homeassistant/helpers/update_coordinator.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ def __init__(
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.
Expand Down Expand Up @@ -141,6 +142,12 @@ def async_update_listeners(self) -> None:
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."""
Expand Down Expand Up @@ -237,7 +244,7 @@ async def _async_refresh( # noqa: C901
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):
Expand Down
37 changes: 37 additions & 0 deletions tests/helpers/test_update_coordinator.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,43 @@ def update_callback():
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:
Expand Down

0 comments on commit ae0cbff

Please sign in to comment.