Replace custom OpenUV data object with coordinators (#80705)

* Replace custom OpenUV data object with coordinators

* Typing

* Code review

homeassistant/components/openuv/__init__.py

@@ -2,11 +2,9 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable
 from typing import Any
 
 from pyopenuv import Client
-from pyopenuv.errors import OpenUvError
 import voluptuous as vol
 
 from homeassistant.config_entries import ConfigEntry, ConfigEntryState
@@ -20,20 +18,16 @@ from homeassistant.const import (
     Platform,
 )
 from homeassistant.core import HomeAssistant, ServiceCall, callback
-from homeassistant.exceptions import ConfigEntryNotReady, HomeAssistantError
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import (
     aiohttp_client,
     config_validation as cv,
     entity_registry,
 )
-from homeassistant.helpers.debounce import Debouncer
+from homeassistant.helpers.entity import EntityDescription
-from homeassistant.helpers.dispatcher import (
-    async_dispatcher_connect,
-    async_dispatcher_send,
-)
-from homeassistant.helpers.entity import Entity, EntityDescription
 from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
 from homeassistant.helpers.service import verify_domain_control
+from homeassistant.helpers.update_coordinator import CoordinatorEntity, UpdateFailed
 
 from .const import (
     CONF_FROM_WINDOW,
@@ -45,13 +39,10 @@ from .const import (
     DOMAIN,
     LOGGER,
 )
+from .coordinator import OpenUvCoordinator
 
 CONF_ENTRY_ID = "entry_id"
 
-DEFAULT_DEBOUNCER_COOLDOWN_SECONDS = 15 * 60
-
-TOPIC_UPDATE = f"{DOMAIN}_data_update"
-
 PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR]
 
 SERVICE_NAME_UPDATE_DATA = "update_data"
@@ -127,53 +118,50 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     _verify_domain_control = verify_domain_control(hass, DOMAIN)
 
     websession = aiohttp_client.async_get_clientsession(hass)
-    openuv = OpenUV(
+    client = Client(
-        hass,
+        entry.data[CONF_API_KEY],
-        entry,
+        entry.data.get(CONF_LATITUDE, hass.config.latitude),
-        Client(
+        entry.data.get(CONF_LONGITUDE, hass.config.longitude),
-            entry.data[CONF_API_KEY],
+        altitude=entry.data.get(CONF_ELEVATION, hass.config.elevation),
-            entry.data.get(CONF_LATITUDE, hass.config.latitude),
+        session=websession,
-            entry.data.get(CONF_LONGITUDE, hass.config.longitude),
-            altitude=entry.data.get(CONF_ELEVATION, hass.config.elevation),
-            session=websession,
-        ),
     )
 
+    async def async_update_protection_data() -> dict[str, Any]:
+        """Update binary sensor (protection window) data."""
+        low = entry.options.get(CONF_FROM_WINDOW, DEFAULT_FROM_WINDOW)
+        high = entry.options.get(CONF_TO_WINDOW, DEFAULT_TO_WINDOW)
+        return await client.uv_protection_window(low=low, high=high)
+
+    coordinators: dict[str, OpenUvCoordinator] = {
+        coordinator_name: OpenUvCoordinator(
+            hass,
+            name=coordinator_name,
+            latitude=client.latitude,
+            longitude=client.longitude,
+            update_method=update_method,
+        )
+        for coordinator_name, update_method in (
+            (DATA_UV, client.uv_index),
+            (DATA_PROTECTION_WINDOW, async_update_protection_data),
+        )
+    }
+
     # We disable the client's request retry abilities here to avoid a lengthy (and
-    # blocking) startup:
+    # blocking) startup; then, if the initial update is successful, we re-enable client
-    openuv.client.disable_request_retries()
+    # request retries:
-
+    client.disable_request_retries()
-    try:
+    init_tasks = [
-        await openuv.async_update()
+        coordinator.async_config_entry_first_refresh()
-    except HomeAssistantError as err:
+        for coordinator in coordinators.values()
-        LOGGER.error("Config entry failed: %s", err)
+    ]
-        raise ConfigEntryNotReady from err
+    await asyncio.gather(*init_tasks)
-
+    client.enable_request_retries()
-    # Once we've successfully authenticated, we re-enable client request retries:
-    openuv.client.enable_request_retries()
 
     hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = openuv
+    hass.data[DOMAIN][entry.entry_id] = coordinators
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
-    @callback
-    def extract_openuv(func: Callable) -> Callable:
-        """Define a decorator to get the correct OpenUV object for a service call."""
-
-        async def wrapper(call: ServiceCall) -> None:
-            """Wrap the service function."""
-            openuv: OpenUV = hass.data[DOMAIN][call.data[CONF_ENTRY_ID]]
-
-            try:
-                await func(call, openuv)
-            except OpenUvError as err:
-                raise HomeAssistantError(
-                    f'Error while executing "{call.service}": {err}'
-                ) from err
-
-        return wrapper
-
     # We determine entity IDs needed to help the user migrate from deprecated services:
     current_uv_index_entity_id = async_get_entity_id_from_unique_id_suffix(
         hass, entry, "current_uv_index"
@@ -183,8 +171,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
 
     @_verify_domain_control
-    @extract_openuv
+    async def update_data(call: ServiceCall) -> None:
-    async def update_data(call: ServiceCall, openuv: OpenUV) -> None:
         """Refresh all OpenUV data."""
         LOGGER.debug("Refreshing all OpenUV data")
         async_log_deprecated_service_call(
@@ -194,12 +181,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             [protection_window_entity_id, current_uv_index_entity_id],
             "2022.12.0",
         )
-        await openuv.async_update()
+
-        async_dispatcher_send(hass, TOPIC_UPDATE)
+        tasks = [coordinator.async_refresh() for coordinator in coordinators.values()]
+        try:
+            await asyncio.gather(*tasks)
+        except UpdateFailed as err:
+            raise HomeAssistantError(err) from err
 
     @_verify_domain_control
-    @extract_openuv
+    async def update_uv_index_data(call: ServiceCall) -> None:
-    async def update_uv_index_data(call: ServiceCall, openuv: OpenUV) -> None:
         """Refresh OpenUV UV index data."""
         LOGGER.debug("Refreshing OpenUV UV index data")
         async_log_deprecated_service_call(
@@ -209,12 +199,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             [current_uv_index_entity_id],
             "2022.12.0",
         )
-        await openuv.async_update_uv_index_data()
+
-        async_dispatcher_send(hass, TOPIC_UPDATE)
+        try:
+            await coordinators[DATA_UV].async_request_refresh()
+        except UpdateFailed as err:
+            raise HomeAssistantError(err) from err
 
     @_verify_domain_control
-    @extract_openuv
+    async def update_protection_data(call: ServiceCall) -> None:
-    async def update_protection_data(call: ServiceCall, openuv: OpenUV) -> None:
         """Refresh OpenUV protection window data."""
         LOGGER.debug("Refreshing OpenUV protection window data")
         async_log_deprecated_service_call(
@@ -224,8 +216,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             [protection_window_entity_id],
             "2022.12.0",
         )
-        await openuv.async_update_protection_data()
+
-        async_dispatcher_send(hass, TOPIC_UPDATE)
+        try:
+            await coordinators[DATA_PROTECTION_WINDOW].async_request_refresh()
+        except UpdateFailed as err:
+            raise HomeAssistantError(err) from err
 
     service_schema = vol.Schema(
         {
@@ -283,106 +278,42 @@ async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-class OpenUV:
+class OpenUvEntity(CoordinatorEntity):
-    """Define a generic OpenUV object."""
-
-    def __init__(self, hass: HomeAssistant, entry: ConfigEntry, client: Client) -> None:
-        """Initialize."""
-        self._update_protection_data_debouncer = Debouncer(
-            hass,
-            LOGGER,
-            cooldown=DEFAULT_DEBOUNCER_COOLDOWN_SECONDS,
-            immediate=True,
-            function=self._async_update_protection_data,
-        )
-
-        self._update_uv_index_data_debouncer = Debouncer(
-            hass,
-            LOGGER,
-            cooldown=DEFAULT_DEBOUNCER_COOLDOWN_SECONDS,
-            immediate=True,
-            function=self._async_update_uv_index_data,
-        )
-
-        self._entry = entry
-        self.client = client
-        self.data: dict[str, Any] = {DATA_PROTECTION_WINDOW: {}, DATA_UV: {}}
-
-    async def _async_update_protection_data(self) -> None:
-        """Update binary sensor (protection window) data."""
-        low = self._entry.options.get(CONF_FROM_WINDOW, DEFAULT_FROM_WINDOW)
-        high = self._entry.options.get(CONF_TO_WINDOW, DEFAULT_TO_WINDOW)
-
-        try:
-            data = await self.client.uv_protection_window(low=low, high=high)
-        except OpenUvError as err:
-            raise HomeAssistantError(
-                f"Error during protection data update: {err}"
-            ) from err
-
-        self.data[DATA_PROTECTION_WINDOW] = data.get("result")
-
-    async def _async_update_uv_index_data(self) -> None:
-        """Update sensor (uv index, etc) data."""
-        try:
-            data = await self.client.uv_index()
-        except OpenUvError as err:
-            raise HomeAssistantError(
-                f"Error during UV index data update: {err}"
-            ) from err
-
-        self.data[DATA_UV] = data.get("result")
-
-    async def async_update_protection_data(self) -> None:
-        """Update binary sensor (protection window) data with a debouncer."""
-        await self._update_protection_data_debouncer.async_call()
-
-    async def async_update_uv_index_data(self) -> None:
-        """Update sensor (uv index, etc) data with a debouncer."""
-        await self._update_uv_index_data_debouncer.async_call()
-
-    async def async_update(self) -> None:
-        """Update sensor/binary sensor data."""
-        tasks = [self.async_update_protection_data(), self.async_update_uv_index_data()]
-        await asyncio.gather(*tasks)
-
-
-class OpenUvEntity(Entity):
     """Define a generic OpenUV entity."""
 
     _attr_has_entity_name = True
 
-    def __init__(self, openuv: OpenUV, description: EntityDescription) -> None:
+    def __init__(
+        self, coordinator: OpenUvCoordinator, description: EntityDescription
+    ) -> None:
         """Initialize."""
+        super().__init__(coordinator)
+
         self._attr_extra_state_attributes = {}
-        self._attr_should_poll = False
         self._attr_unique_id = (
-            f"{openuv.client.latitude}_{openuv.client.longitude}_{description.key}"
+            f"{coordinator.latitude}_{coordinator.longitude}_{description.key}"
         )
         self.entity_description = description
-        self.openuv = openuv
 
     @callback
-    def async_update_state(self) -> None:
+    def _handle_coordinator_update(self) -> None:
-        """Update the state."""
+        """Respond to a DataUpdateCoordinator update."""
-        self.update_from_latest_data()
+        self._update_from_latest_data()
         self.async_write_ha_state()
 
+    @callback
+    def _update_from_latest_data(self) -> None:
+        """Update the entity from the latest data."""
+        raise NotImplementedError
+
     async def async_added_to_hass(self) -> None:
-        """Register callbacks."""
+        """Handle entity which will be added."""
-        self.update_from_latest_data()
+        await super().async_added_to_hass()
-        self.async_on_remove(
+        self._update_from_latest_data()
-            async_dispatcher_connect(self.hass, TOPIC_UPDATE, self.async_update_state)
-        )
 
     async def async_update(self) -> None:
         """Update the entity.
 
-        Only used by the generic entity update service. Should be implemented by each
+        Only used by the generic entity update service.
-        OpenUV platform.
         """
-        raise NotImplementedError
+        await self.coordinator.async_request_refresh()
-
-    def update_from_latest_data(self) -> None:
-        """Update the sensor using the latest data."""
-        raise NotImplementedError

homeassistant/components/openuv/binary_sensor.py

@@ -10,6 +10,7 @@ from homeassistant.util.dt import as_local, parse_datetime, utcnow
 
 from . import OpenUvEntity
 from .const import DATA_PROTECTION_WINDOW, DOMAIN, LOGGER, TYPE_PROTECTION_WINDOW
+from .coordinator import OpenUvCoordinator
 
 ATTR_PROTECTION_WINDOW_ENDING_TIME = "end_time"
 ATTR_PROTECTION_WINDOW_ENDING_UV = "end_uv"
@@ -26,32 +27,27 @@ BINARY_SENSOR_DESCRIPTION_PROTECTION_WINDOW = BinarySensorEntityDescription(
 async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
+    # Once we've successfully authenticated, we re-enable client request retries:
     """Set up an OpenUV sensor based on a config entry."""
-    openuv = hass.data[DOMAIN][entry.entry_id]
+    coordinators: dict[str, OpenUvCoordinator] = hass.data[DOMAIN][entry.entry_id]
+
     async_add_entities(
-        [OpenUvBinarySensor(openuv, BINARY_SENSOR_DESCRIPTION_PROTECTION_WINDOW)]
+        [
+            OpenUvBinarySensor(
+                coordinators[DATA_PROTECTION_WINDOW],
+                BINARY_SENSOR_DESCRIPTION_PROTECTION_WINDOW,
+            )
+        ]
     )
 
 
 class OpenUvBinarySensor(OpenUvEntity, BinarySensorEntity):
     """Define a binary sensor for OpenUV."""
 
-    async def async_update(self) -> None:
-        """Update the entity.
-
-        Only used by the generic entity update service.
-        """
-        await self.openuv.async_update_protection_data()
-        self.async_update_state()
-
     @callback
-    def update_from_latest_data(self) -> None:
+    def _update_from_latest_data(self) -> None:
-        """Update the state."""
+        """Update the entity from the latest data."""
-        if not (data := self.openuv.data[DATA_PROTECTION_WINDOW]):
+        data = self.coordinator.data
-            self._attr_available = False
-            return
-
-        self._attr_available = True
 
         for key in ("from_time", "to_time", "from_uv", "to_uv"):
             if not data.get(key):

homeassistant/components/openuv/coordinator.py

@@ -0,0 +1,55 @@
+"""Define an update coordinator for OpenUV."""
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any, cast
+
+from pyopenuv.errors import OpenUvError
+
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.debounce import Debouncer
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import LOGGER
+
+DEFAULT_DEBOUNCER_COOLDOWN_SECONDS = 15 * 60
+
+
+class OpenUvCoordinator(DataUpdateCoordinator):
+    """Define an OpenUV data coordinator."""
+
+    update_method: Callable[[], Awaitable[dict[str, Any]]]
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        *,
+        name: str,
+        latitude: str,
+        longitude: str,
+        update_method: Callable[[], Awaitable[dict[str, Any]]],
+    ) -> None:
+        """Initialize."""
+        super().__init__(
+            hass,
+            LOGGER,
+            name=name,
+            update_method=update_method,
+            request_refresh_debouncer=Debouncer(
+                hass,
+                LOGGER,
+                cooldown=DEFAULT_DEBOUNCER_COOLDOWN_SECONDS,
+                immediate=True,
+            ),
+        )
+
+        self.latitude = latitude
+        self.longitude = longitude
+
+    async def _async_update_data(self) -> dict[str, Any]:
+        """Fetch data from OpenUV."""
+        try:
+            data = await self.update_method()
+        except OpenUvError as err:
+            raise UpdateFailed(f"Error during protection data update: {err}") from err
+        return cast(dict[str, Any], data["result"])

homeassistant/components/openuv/diagnostics.py

@@ -13,8 +13,8 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant
 
-from . import OpenUV
 from .const import DOMAIN
+from .coordinator import OpenUvCoordinator
 
 CONF_COORDINATES = "coordinates"
 CONF_TITLE = "title"
@@ -33,9 +33,15 @@ async def async_get_config_entry_diagnostics(
     hass: HomeAssistant, entry: ConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    openuv: OpenUV = hass.data[DOMAIN][entry.entry_id]
+    coordinators: dict[str, OpenUvCoordinator] = hass.data[DOMAIN][entry.entry_id]
 
-    return {
+    return async_redact_data(
-        "entry": async_redact_data(entry.as_dict(), TO_REDACT),
+        {
-        "data": async_redact_data(openuv.data, TO_REDACT),
+            "entry": entry.as_dict(),
-    }
+            "data": {
+                coordinator_name: coordinator.data
+                for coordinator_name, coordinator in coordinators.items()
+            },
+        },
+        TO_REDACT,
+    )

homeassistant/components/openuv/sensor.py

@@ -28,6 +28,7 @@ from .const import (
     TYPE_SAFE_EXPOSURE_TIME_5,
     TYPE_SAFE_EXPOSURE_TIME_6,
 )
+from .coordinator import OpenUvCoordinator
 
 ATTR_MAX_UV_TIME = "time"
 
@@ -122,31 +123,23 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up a OpenUV sensor based on a config entry."""
-    openuv = hass.data[DOMAIN][entry.entry_id]
+    coordinators: dict[str, OpenUvCoordinator] = hass.data[DOMAIN][entry.entry_id]
+
     async_add_entities(
-        [OpenUvSensor(openuv, description) for description in SENSOR_DESCRIPTIONS]
+        [
+            OpenUvSensor(coordinators[DATA_UV], description)
+            for description in SENSOR_DESCRIPTIONS
+        ]
     )
 
 
 class OpenUvSensor(OpenUvEntity, SensorEntity):
     """Define a binary sensor for OpenUV."""
 
-    async def async_update(self) -> None:
-        """Update the entity.
-
-        Only used by the generic entity update service.
-        """
-        await self.openuv.async_update_uv_index_data()
-        self.async_update_state()
-
     @callback
-    def update_from_latest_data(self) -> None:
+    def _update_from_latest_data(self) -> None:
         """Update the state."""
-        if (data := self.openuv.data[DATA_UV]) is None:
+        data = self.coordinator.data
-            self._attr_available = False
-            return
-
-        self._attr_available = True
 
         if self.entity_description.key == TYPE_CURRENT_OZONE_LEVEL:
             self._attr_native_value = data["ozone"]

tests/components/openuv/test_diagnostics.py

@@ -1,6 +1,5 @@
 """Test OpenUV diagnostics."""
 from homeassistant.components.diagnostics import REDACTED
-from homeassistant.const import CONF_ENTITY_ID
 from homeassistant.setup import async_setup_component
 
 from tests.components.diagnostics import get_diagnostics_for_config_entry
@@ -9,12 +8,6 @@ from tests.components.diagnostics import get_diagnostics_for_config_entry
 async def test_entry_diagnostics(hass, config_entry, hass_client, setup_openuv):
     """Test config entry diagnostics."""
     await async_setup_component(hass, "homeassistant", {})
-    await hass.services.async_call(
-        "homeassistant",
-        "update_entity",
-        {CONF_ENTITY_ID: ["sensor.current_uv_index"]},
-        blocking=True,
-    )
     assert await get_diagnostics_for_config_entry(hass, hass_client, config_entry) == {
         "entry": {
             "entry_id": config_entry.entry_id,