Add update class method to DataUpdateCoordinator

homeassistant/helpers/update_coordinator.py

@@ -30,8 +30,8 @@ def __init__(
         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 @@ async def async_request_refresh(self) -> None:
         """
         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 @@ async def async_refresh(self) -> None:
 
         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 @@ async def async_refresh(self) -> None:
                 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(

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