Convert solax to use DataUpdateCoordinator (#117767)

homeassistant/components/solax/__init__.py

@@ -1,18 +1,39 @@
 """The solax component."""
 
-from solax import real_time_api
+from dataclasses import dataclass
+from datetime import timedelta
+import logging
+
+from solax import InverterResponse, RealTimeAPI, real_time_api
+from solax.inverter import InverterError
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_IP_ADDRESS, CONF_PASSWORD, CONF_PORT, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.helpers.update_coordinator import UpdateFailed
 
-from .const import DOMAIN
+from .coordinator import SolaxDataUpdateCoordinator
 
 PLATFORMS = [Platform.SENSOR]
 
+SCAN_INTERVAL = timedelta(seconds=30)
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+@dataclass(slots=True)
+class SolaxData:
+    """Class for storing solax data."""
+
+    api: RealTimeAPI
+    coordinator: SolaxDataUpdateCoordinator
+
+
+type SolaxConfigEntry = ConfigEntry[SolaxData]
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: SolaxConfigEntry) -> bool:
     """Set up the sensors from a ConfigEntry."""
 
     try:
@@ -21,19 +42,30 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             entry.data[CONF_PORT],
             entry.data[CONF_PASSWORD],
         )
-        await api.get_data()
     except Exception as err:
         raise ConfigEntryNotReady from err
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = api
+    async def _async_update() -> InverterResponse:
+        try:
+            return await api.get_data()
+        except InverterError as err:
+            raise UpdateFailed from err
+
+    coordinator = SolaxDataUpdateCoordinator(
+        hass,
+        logger=_LOGGER,
+        name=f"solax {entry.title}",
+        update_interval=SCAN_INTERVAL,
+        update_method=_async_update,
+    )
+    await coordinator.async_config_entry_first_refresh()
+
+    entry.runtime_data = SolaxData(api=api, coordinator=coordinator)
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: SolaxConfigEntry) -> bool:
     """Unload a config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok

homeassistant/components/solax/coordinator.py

@@ -0,0 +1,9 @@
+"""Constants for the solax integration."""
+
+from solax import InverterResponse
+
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
+
+
+class SolaxDataUpdateCoordinator(DataUpdateCoordinator[InverterResponse]):
+    """DataUpdateCoordinator for solax."""

homeassistant/components/solax/sensor.py

@@ -2,11 +2,6 @@
 
 from __future__ import annotations
 
-import asyncio
-from datetime import timedelta
-
-from solax import RealTimeAPI
-from solax.inverter import InverterError
 from solax.units import Units
 
 from homeassistant.components.sensor import (
@@ -15,7 +10,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     PERCENTAGE,
     UnitOfElectricCurrent,
@@ -26,15 +20,15 @@ from homeassistant.const import (
     UnitOfTemperature,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import PlatformNotReady
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.event import async_track_time_interval
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from . import SolaxConfigEntry
 from .const import DOMAIN, MANUFACTURER
+from .coordinator import SolaxDataUpdateCoordinator
 
 DEFAULT_PORT = 80
-SCAN_INTERVAL = timedelta(seconds=30)
 
 
 SENSOR_DESCRIPTIONS: dict[tuple[Units, bool], SensorEntityDescription] = {
@@ -94,28 +88,23 @@ SENSOR_DESCRIPTIONS: dict[tuple[Units, bool], SensorEntityDescription] = {
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: SolaxConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Entry setup."""
-    api: RealTimeAPI = hass.data[DOMAIN][entry.entry_id]
+    api = entry.runtime_data.api
-    resp = await api.get_data()
+    coordinator = entry.runtime_data.coordinator
+    resp = coordinator.data
     serial = resp.serial_number
     version = resp.version
-    endpoint = RealTimeDataEndpoint(hass, api)
+    entities: list[InverterSensorEntity] = []
-    entry.async_create_background_task(
-        hass, endpoint.async_refresh(), f"solax {entry.title} initial refresh"
-    )
-    entry.async_on_unload(
-        async_track_time_interval(hass, endpoint.async_refresh, SCAN_INTERVAL)
-    )
-    devices = []
     for sensor, (idx, measurement) in api.inverter.sensor_map().items():
         description = SENSOR_DESCRIPTIONS[(measurement.unit, measurement.is_monotonic)]
 
         uid = f"{serial}-{idx}"
-        devices.append(
+        entities.append(
-            Inverter(
+            InverterSensorEntity(
+                coordinator,
                 api.inverter.manufacturer,
                 uid,
                 serial,
@@ -126,57 +115,28 @@ async def async_setup_entry(
                 description.device_class,
             )
         )
-    endpoint.sensors = devices
+    async_add_entities(entities)
-    async_add_entities(devices)
 
 
-class RealTimeDataEndpoint:
+class InverterSensorEntity(CoordinatorEntity, SensorEntity):
-    """Representation of a Sensor."""
-
-    def __init__(self, hass: HomeAssistant, api: RealTimeAPI) -> None:
-        """Initialize the sensor."""
-        self.hass = hass
-        self.api = api
-        self.ready = asyncio.Event()
-        self.sensors: list[Inverter] = []
-
-    async def async_refresh(self, now=None):
-        """Fetch new state data for the sensor.
-
-        This is the only method that should fetch new data for Home Assistant.
-        """
-        try:
-            api_response = await self.api.get_data()
-            self.ready.set()
-        except InverterError as err:
-            if now is not None:
-                self.ready.clear()
-                return
-            raise PlatformNotReady from err
-        data = api_response.data
-        for sensor in self.sensors:
-            if sensor.key in data:
-                sensor.value = data[sensor.key]
-                sensor.async_schedule_update_ha_state()
-
-
-class Inverter(SensorEntity):
     """Class for a sensor."""
 
     _attr_should_poll = False
 
     def __init__(
         self,
-        manufacturer,
+        coordinator: SolaxDataUpdateCoordinator,
-        uid,
+        manufacturer: str,
-        serial,
+        uid: str,
-        version,
+        serial: str,
-        key,
+        version: str,
-        unit,
+        key: str,
-        state_class=None,
+        unit: str | None,
-        device_class=None,
+        state_class: SensorStateClass | str | None,
-    ):
+        device_class: SensorDeviceClass | None,
+    ) -> None:
         """Initialize an inverter sensor."""
+        super().__init__(coordinator)
         self._attr_unique_id = uid
         self._attr_name = f"{manufacturer} {serial} {key}"
         self._attr_native_unit_of_measurement = unit
@@ -189,9 +149,8 @@ class Inverter(SensorEntity):
             sw_version=version,
         )
         self.key = key
-        self.value = None
 
     @property
     def native_value(self):
         """State of this inverter attribute."""
-        return self.value
+        return self.coordinator.data.data[self.key]