From f8712b0e00fc3d27edab928b3a38190389461338 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 30 Aug 2020 05:02:37 -0500
Subject: [PATCH] Create a CoordinatorEntity class to avoid repating code in
 integrations (#39388)

---
 homeassistant/helpers/update_coordinator.py | 33 ++++++++++++++++++++-
 tests/helpers/test_update_coordinator.py    | 22 +++++++++++++-
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/homeassistant/helpers/update_coordinator.py b/homeassistant/helpers/update_coordinator.py
index 7b7e6af4d62..987f1d63eee 100644
--- a/homeassistant/helpers/update_coordinator.py
+++ b/homeassistant/helpers/update_coordinator.py
@@ -10,7 +10,7 @@ import aiohttp
 import requests
 
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
-from homeassistant.helpers import event
+from homeassistant.helpers import entity, event
 from homeassistant.util.dt import utcnow
 
 from .debounce import Debouncer
@@ -190,3 +190,34 @@ class DataUpdateCoordinator(Generic[T]):
 
         for update_callback in self._listeners:
             update_callback()
+
+
+class CoordinatorEntity(entity.Entity):
+    """A class for entities using DataUpdateCoordinator."""
+
+    def __init__(self, coordinator: DataUpdateCoordinator) -> None:
+        """Create the entity with a DataUpdateCoordinator."""
+        self.coordinator = coordinator
+
+    @property
+    def should_poll(self) -> bool:
+        """No need to poll. Coordinator notifies entity of updates."""
+        return False
+
+    @property
+    def available(self) -> bool:
+        """Return if entity is available."""
+        return self.coordinator.last_update_success
+
+    async def async_added_to_hass(self) -> None:
+        """When entity is added to hass."""
+        self.async_on_remove(
+            self.coordinator.async_add_listener(self.async_write_ha_state)
+        )
+
+    async def async_update(self) -> None:
+        """Update the entity.
+
+        Only used by the generic entity update service.
+        """
+        await self.coordinator.async_request_refresh()
diff --git a/tests/helpers/test_update_coordinator.py b/tests/helpers/test_update_coordinator.py
index 56c53f1994c..73360a3053b 100644
--- a/tests/helpers/test_update_coordinator.py
+++ b/tests/helpers/test_update_coordinator.py
@@ -11,7 +11,7 @@ import requests
 from homeassistant.helpers import update_coordinator
 from homeassistant.util.dt import utcnow
 
-from tests.async_mock import AsyncMock, Mock
+from tests.async_mock import AsyncMock, Mock, patch
 from tests.common import async_fire_time_changed
 
 LOGGER = logging.getLogger(__name__)
@@ -224,3 +224,23 @@ async def test_refresh_recover(crd, caplog):
 
     assert crd.last_update_success is True
     assert "Fetching test data recovered" in caplog.text
+
+
+async def test_coordinator_entity(crd):
+    """Test the CoordinatorEntity class."""
+    entity = update_coordinator.CoordinatorEntity(crd)
+
+    assert entity.should_poll is False
+
+    crd.last_update_success = False
+    assert entity.available is False
+
+    await entity.async_update()
+    assert entity.available is True
+
+    with patch(
+        "homeassistant.helpers.entity.Entity.async_on_remove"
+    ) as mock_async_on_remove:
+        await entity.async_added_to_hass()
+
+    assert mock_async_on_remove.called