Add support for manual updating a data coordinator (#42746)

This commit is contained in:
Paulus Schoutsen 2020-11-02 09:54:08 +01:00 committed by GitHub
parent ce16ff27ef
commit 1844e30858
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 57 additions and 1 deletions

View file

@ -138,9 +138,9 @@ class DataUpdateCoordinator(Generic[T]):
self._unsub_refresh = None
self._debounced_refresh.async_cancel()
start = monotonic()
try:
start = monotonic()
self.data = await self._async_update_data()
except (asyncio.TimeoutError, requests.exceptions.Timeout):
@ -192,6 +192,28 @@ class DataUpdateCoordinator(Generic[T]):
for update_callback in self._listeners:
update_callback()
@callback
def async_set_updated_data(self, data: T) -> None:
"""Manually update data, notify listeners and reset refresh interval."""
if self._unsub_refresh:
self._unsub_refresh()
self._unsub_refresh = None
self._debounced_refresh.async_cancel()
self.data = data
self.last_update_success = True
self.logger.debug(
"Manually updated %s data",
self.name,
)
if self._listeners:
self._schedule_refresh()
for update_callback in self._listeners:
update_callback()
class CoordinatorEntity(entity.Entity):
"""A class for entities using DataUpdateCoordinator."""

View file

@ -250,3 +250,37 @@ async def test_coordinator_entity(crd):
with patch("homeassistant.helpers.entity.Entity.enabled", False):
await entity.async_update()
assert entity.available is False
async def test_async_set_updated_data(crd):
"""Test async_set_updated_data for update coordinator."""
assert crd.data is None
with patch.object(crd._debounced_refresh, "async_cancel") as mock_cancel:
crd.async_set_updated_data(100)
# Test we cancel any pending refresh
assert len(mock_cancel.mock_calls) == 1
# Test data got updated
assert crd.data == 100
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)
crd.async_set_updated_data(200)
assert updates == [200]
assert crd._unsub_refresh is not None
old_refresh = crd._unsub_refresh
crd.async_set_updated_data(300)
# We have created a new refresh listener
assert crd._unsub_refresh is not old_refresh