From 1844e30858abdb6dc4121808b42db16e2d046cfa Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <balloob@gmail.com>
Date: Mon, 2 Nov 2020 09:54:08 +0100
Subject: [PATCH] Add support for manual updating a data coordinator (#42746)

---
 homeassistant/helpers/update_coordinator.py | 24 ++++++++++++++-
 tests/helpers/test_update_coordinator.py    | 34 +++++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/homeassistant/helpers/update_coordinator.py b/homeassistant/helpers/update_coordinator.py
index 2abe7b8c6b7..895eff01f57 100644
--- a/homeassistant/helpers/update_coordinator.py
+++ b/homeassistant/helpers/update_coordinator.py
@@ -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."""
diff --git a/tests/helpers/test_update_coordinator.py b/tests/helpers/test_update_coordinator.py
index 72b9bff60f1..90567456bfb 100644
--- a/tests/helpers/test_update_coordinator.py
+++ b/tests/helpers/test_update_coordinator.py
@@ -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