From ff92a8b2602f38f373a7ed4954626fc08d58c084 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Fri, 13 Mar 2020 00:27:19 +0100
Subject: [PATCH] Add update class method to DataUpdateCoordinator (#32724)

* Add update class method to DataUpdateCoordinator

* Update homeassistant/helpers/update_coordinator.py

Co-Authored-By: Paulus Schoutsen <balloob@gmail.com>

* Move update_method param

* Rename async_update_data to _async_update_data

* Raise NotImplementedError

* Re-raise NotImplementedError

* Remove caplog, not needed anymore

* Don't set last_update_success on NotImplementedError

* Fix test

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>
---
 homeassistant/helpers/update_coordinator.py | 15 ++++++++++++---
 tests/helpers/test_update_coordinator.py    | 20 ++++++++++++++++++++
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/homeassistant/helpers/update_coordinator.py b/homeassistant/helpers/update_coordinator.py
index b2fe87148b1..85db657c441 100644
--- a/homeassistant/helpers/update_coordinator.py
+++ b/homeassistant/helpers/update_coordinator.py
@@ -30,8 +30,8 @@ class DataUpdateCoordinator:
         logger: logging.Logger,
         *,
         name: str,
-        update_method: Callable[[], Awaitable],
         update_interval: timedelta,
+        update_method: Optional[Callable[[], Awaitable]] = None,
         request_refresh_debouncer: Optional[Debouncer] = None,
     ):
         """Initialize global data updater."""
@@ -104,8 +104,14 @@ class DataUpdateCoordinator:
         """
         await self._debounced_refresh.async_call()
 
+    async def _async_update_data(self) -> Optional[Any]:
+        """Fetch the latest data from the source."""
+        if self.update_method is None:
+            raise NotImplementedError("Update method not implemented")
+        return await self.update_method()
+
     async def async_refresh(self) -> None:
-        """Update data."""
+        """Refresh data."""
         if self._unsub_refresh:
             self._unsub_refresh()
             self._unsub_refresh = None
@@ -114,7 +120,7 @@ class DataUpdateCoordinator:
 
         try:
             start = monotonic()
-            self.data = await self.update_method()
+            self.data = await self._async_update_data()
 
         except asyncio.TimeoutError:
             if self.last_update_success:
@@ -131,6 +137,9 @@ class DataUpdateCoordinator:
                 self.logger.error("Error fetching %s data: %s", self.name, err)
                 self.last_update_success = False
 
+        except NotImplementedError as err:
+            raise err
+
         except Exception as err:  # pylint: disable=broad-except
             self.last_update_success = False
             self.logger.exception(
diff --git a/tests/helpers/test_update_coordinator.py b/tests/helpers/test_update_coordinator.py
index 115e00168fc..c17c79ccbc8 100644
--- a/tests/helpers/test_update_coordinator.py
+++ b/tests/helpers/test_update_coordinator.py
@@ -104,6 +104,16 @@ async def test_refresh_fail_unknown(crd, caplog):
     assert "Unexpected error fetching test data" in caplog.text
 
 
+async def test_refresh_no_update_method(crd):
+    """Test raising error is no update method is provided."""
+    await crd.async_refresh()
+
+    crd.update_method = None
+
+    with pytest.raises(NotImplementedError):
+        await crd.async_refresh()
+
+
 async def test_update_interval(hass, crd):
     """Test update interval works."""
     # Test we don't update without subscriber
@@ -132,3 +142,13 @@ async def test_update_interval(hass, crd):
 
     # Test we stop updating after we lose last subscriber
     assert crd.data == 2
+
+
+async def test_refresh_recover(crd, caplog):
+    """Test recovery of freshing data."""
+    crd.last_update_success = False
+
+    await crd.async_refresh()
+
+    assert crd.last_update_success is True
+    assert "Fetching test data recovered" in caplog.text