Add update class method to DataUpdateCoordinator

homeassistant/helpers/update_coordinator.py

@@ -30,7 +30,7 @@ def __init__(
         logger: logging.Logger,
         *,
         name: str,
-        update_method: Callable[[], Awaitable],
+        update_method: Optional[Callable[[], Awaitable]] = None,
         update_interval: timedelta,
         request_refresh_debouncer: Optional[Debouncer] = None,
     ):
@@ -104,8 +104,14 @@ async def async_request_refresh(self) -> None:
         """
         await self._debounced_refresh.async_call()
 
-    async def async_refresh(self) -> None:
+    async def async_update_data(self) -> Optional[Any]:
         """Update data."""
+        if self.update_method is None:
+            raise UpdateFailed("Update method not implemented")
+        return await self.update_method()
+
+    async def async_refresh(self) -> None:
+        """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:

tests/helpers/test_update_coordinator.py

@@ -104,6 +104,18 @@ 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, caplog):
+    """Test raising error is no update method is provided."""
+    await crd.async_refresh()
+
+    crd.update_method = None
+
+    await crd.async_refresh()
+
+    assert crd.last_update_success is False
+    assert "Update method not implemented" in caplog.text
+
+
 async def test_update_interval(hass, crd):
     """Test update interval works."""
     # Test we don't update without subscriber
@@ -132,3 +144,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