Implement data coordinator for Adax-integration (#139514)

* Implemented coordinator (for Cloud integration)

* Optimized coordinator updates

* Finalizing

* Running ruff and ruff format

* Raise error if trying to instantiate coordinator for a AdaxLocal config

* Re-added data-handler for AdaxLocal integrations

* Added a coordinator for Local integrations

* mypy warnings

* Update homeassistant/components/adax/manifest.json

Co-authored-by: Daniel Hjelseth Høyer <mail@dahoiv.net>

* Resolve mypy issues

* PR comments

- Explicit passing of config_entry to Coordinator base type
- Avoid duplicate storing of Coordinator data. Instead use self.data
- Remove try-catch and wrapping to UpdateFailed in _async_update_data
- Custom ConfigEntry type for passing coordinator via entry.runtime_data

* When changing HVAC_MODE update data via Coordinator to optimize

* Apply already loaded data for Climate entity directly in __init__

* Moved SCAN_INTERVAL into const.py

* Removed logging statements

* Remove unnecessary get_rooms() / get_status() functions

* Resolvning mypy issues

* Adding tests for coordinators

* Resolving review comments by joostlek

* Setup of Cloud devices with device_id

* Implement Climate tests for Adax

* Implementing assertions of UNAVAILABLE state

* Removed no longer needed method

* Apply suggestions from code review

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

* Mock Adax class instead of individual methods

* Mock config entry via fixture

* Load config entry data from .json fixture

* Hard code config_entry_data instead of .json file

* Removed obsolete .json-files

* Fix

* Fix

---------

Co-authored-by: Daniel Hjelseth Høyer <mail@dahoiv.net>
Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

homeassistant/components/adax/__init__.py

@@ -2,25 +2,38 @@
 
 from __future__ import annotations
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 
+from .const import CONNECTION_TYPE, LOCAL
+from .coordinator import AdaxCloudCoordinator, AdaxConfigEntry, AdaxLocalCoordinator
+
 PLATFORMS = [Platform.CLIMATE]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: AdaxConfigEntry) -> bool:
     """Set up Adax from a config entry."""
+    if entry.data.get(CONNECTION_TYPE) == LOCAL:
+        local_coordinator = AdaxLocalCoordinator(hass, entry)
+        entry.runtime_data = local_coordinator
+    else:
+        cloud_coordinator = AdaxCloudCoordinator(hass, entry)
+        entry.runtime_data = cloud_coordinator
+
+    await entry.runtime_data.async_config_entry_first_refresh()
+
     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: AdaxConfigEntry) -> bool:
     """Unload a config entry."""
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
-async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_migrate_entry(
+    hass: HomeAssistant, config_entry: AdaxConfigEntry
+) -> bool:
     """Migrate old entry."""
     # convert title and unique_id to string
     if config_entry.version == 1:

homeassistant/components/adax/climate.py

@@ -12,57 +12,42 @@ from homeassistant.components.climate import (
     ClimateEntityFeature,
     HVACMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_TEMPERATURE,
-    CONF_IP_ADDRESS,
-    CONF_PASSWORD,
-    CONF_TOKEN,
     CONF_UNIQUE_ID,
     PRECISION_WHOLE,
     UnitOfTemperature,
 )
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from .const import ACCOUNT_ID, CONNECTION_TYPE, DOMAIN, LOCAL
+from . import AdaxConfigEntry
+from .const import CONNECTION_TYPE, DOMAIN, LOCAL
+from .coordinator import AdaxCloudCoordinator, AdaxLocalCoordinator
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: AdaxConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the Adax thermostat with config flow."""
     if entry.data.get(CONNECTION_TYPE) == LOCAL:
-        adax_data_handler = AdaxLocal(
+        local_coordinator = cast(AdaxLocalCoordinator, entry.runtime_data)
-            entry.data[CONF_IP_ADDRESS],
-            entry.data[CONF_TOKEN],
-            websession=async_get_clientsession(hass, verify_ssl=False),
-        )
         async_add_entities(
-            [LocalAdaxDevice(adax_data_handler, entry.data[CONF_UNIQUE_ID])], True
+            [LocalAdaxDevice(local_coordinator, entry.data[CONF_UNIQUE_ID])],
+        )
+    else:
+        cloud_coordinator = cast(AdaxCloudCoordinator, entry.runtime_data)
+        async_add_entities(
+            AdaxDevice(cloud_coordinator, device_id)
+            for device_id in cloud_coordinator.data
         )
-        return
-
-    adax_data_handler = Adax(
-        entry.data[ACCOUNT_ID],
-        entry.data[CONF_PASSWORD],
-        websession=async_get_clientsession(hass),
-    )
-
-    async_add_entities(
-        (
-            AdaxDevice(room, adax_data_handler)
-            for room in await adax_data_handler.get_rooms()
-        ),
-        True,
-    )
 
 
-class AdaxDevice(ClimateEntity):
+class AdaxDevice(CoordinatorEntity[AdaxCloudCoordinator], ClimateEntity):
     """Representation of a heater."""
 
     _attr_hvac_modes = [HVACMode.HEAT, HVACMode.OFF]
@@ -76,20 +61,37 @@ class AdaxDevice(ClimateEntity):
     _attr_target_temperature_step = PRECISION_WHOLE
     _attr_temperature_unit = UnitOfTemperature.CELSIUS
 
-    def __init__(self, heater_data: dict[str, Any], adax_data_handler: Adax) -> None:
+    def __init__(
+        self,
+        coordinator: AdaxCloudCoordinator,
+        device_id: str,
+    ) -> None:
         """Initialize the heater."""
-        self._device_id = heater_data["id"]
+        super().__init__(coordinator)
-        self._adax_data_handler = adax_data_handler
+        self._adax_data_handler: Adax = coordinator.adax_data_handler
+        self._device_id = device_id
 
-        self._attr_unique_id = f"{heater_data['homeId']}_{heater_data['id']}"
+        self._attr_name = self.room["name"]
+        self._attr_unique_id = f"{self.room['homeId']}_{self._device_id}"
         self._attr_device_info = DeviceInfo(
-            identifiers={(DOMAIN, heater_data["id"])},
+            identifiers={(DOMAIN, device_id)},
             # Instead of setting the device name to the entity name, adax
             # should be updated to set has_entity_name = True, and set the entity
             # name to None
             name=cast(str | None, self.name),
             manufacturer="Adax",
         )
+        self._apply_data(self.room)
+
+    @property
+    def available(self) -> bool:
+        """Whether the entity is available or not."""
+        return super().available and self._device_id in self.coordinator.data
+
+    @property
+    def room(self) -> dict[str, Any]:
+        """Gets the data for this particular device."""
+        return self.coordinator.data[self._device_id]
 
     async def async_set_hvac_mode(self, hvac_mode: HVACMode) -> None:
         """Set hvac mode."""
@@ -104,7 +106,9 @@ class AdaxDevice(ClimateEntity):
             )
         else:
             return
-        await self._adax_data_handler.update()
+
+        # Request data refresh from source to verify that update was successful
+        await self.coordinator.async_request_refresh()
 
     async def async_set_temperature(self, **kwargs: Any) -> None:
         """Set new target temperature."""
@@ -114,28 +118,31 @@ class AdaxDevice(ClimateEntity):
             self._device_id, temperature, True
         )
 
-    async def async_update(self) -> None:
+    @callback
-        """Get the latest data."""
+    def _handle_coordinator_update(self) -> None:
-        for room in await self._adax_data_handler.get_rooms():
+        """Handle updated data from the coordinator."""
-            if room["id"] != self._device_id:
+        if room := self.room:
-                continue
+            self._apply_data(room)
-            self._attr_name = room["name"]
+        super()._handle_coordinator_update()
-            self._attr_current_temperature = room.get("temperature")
+
-            self._attr_target_temperature = room.get("targetTemperature")
+    def _apply_data(self, room: dict[str, Any]) -> None:
-            if room["heatingEnabled"]:
+        """Update the appropriate attributues based on received data."""
-                self._attr_hvac_mode = HVACMode.HEAT
+        self._attr_current_temperature = room.get("temperature")
-                self._attr_icon = "mdi:radiator"
+        self._attr_target_temperature = room.get("targetTemperature")
-            else:
+        if room["heatingEnabled"]:
-                self._attr_hvac_mode = HVACMode.OFF
+            self._attr_hvac_mode = HVACMode.HEAT
-                self._attr_icon = "mdi:radiator-off"
+            self._attr_icon = "mdi:radiator"
-            return
+        else:
+            self._attr_hvac_mode = HVACMode.OFF
+            self._attr_icon = "mdi:radiator-off"
 
 
-class LocalAdaxDevice(ClimateEntity):
+class LocalAdaxDevice(CoordinatorEntity[AdaxLocalCoordinator], ClimateEntity):
     """Representation of a heater."""
 
     _attr_hvac_modes = [HVACMode.HEAT, HVACMode.OFF]
-    _attr_hvac_mode = HVACMode.HEAT
+    _attr_hvac_mode = HVACMode.OFF
+    _attr_icon = "mdi:radiator-off"
     _attr_max_temp = 35
     _attr_min_temp = 5
     _attr_supported_features = (
@@ -146,9 +153,10 @@ class LocalAdaxDevice(ClimateEntity):
     _attr_target_temperature_step = PRECISION_WHOLE
     _attr_temperature_unit = UnitOfTemperature.CELSIUS
 
-    def __init__(self, adax_data_handler: AdaxLocal, unique_id: str) -> None:
+    def __init__(self, coordinator: AdaxLocalCoordinator, unique_id: str) -> None:
         """Initialize the heater."""
-        self._adax_data_handler = adax_data_handler
+        super().__init__(coordinator)
+        self._adax_data_handler: AdaxLocal = coordinator.adax_data_handler
         self._attr_unique_id = unique_id
         self._attr_device_info = DeviceInfo(
             identifiers={(DOMAIN, unique_id)},
@@ -169,17 +177,20 @@ class LocalAdaxDevice(ClimateEntity):
             return
         await self._adax_data_handler.set_target_temperature(temperature)
 
-    async def async_update(self) -> None:
+    @callback
-        """Get the latest data."""
+    def _handle_coordinator_update(self) -> None:
-        data = await self._adax_data_handler.get_status()
+        """Handle updated data from the coordinator."""
-        self._attr_current_temperature = data["current_temperature"]
+        if data := self.coordinator.data:
-        self._attr_available = self._attr_current_temperature is not None
+            self._attr_current_temperature = data["current_temperature"]
-        if (target_temp := data["target_temperature"]) == 0:
+            self._attr_available = self._attr_current_temperature is not None
-            self._attr_hvac_mode = HVACMode.OFF
+            if (target_temp := data["target_temperature"]) == 0:
-            self._attr_icon = "mdi:radiator-off"
+                self._attr_hvac_mode = HVACMode.OFF
-            if target_temp == 0:
+                self._attr_icon = "mdi:radiator-off"
-                self._attr_target_temperature = self._attr_min_temp
+                if target_temp == 0:
-        else:
+                    self._attr_target_temperature = self._attr_min_temp
-            self._attr_hvac_mode = HVACMode.HEAT
+            else:
-            self._attr_icon = "mdi:radiator"
+                self._attr_hvac_mode = HVACMode.HEAT
-            self._attr_target_temperature = target_temp
+                self._attr_icon = "mdi:radiator"
+                self._attr_target_temperature = target_temp
+
+        super()._handle_coordinator_update()

homeassistant/components/adax/const.py

@@ -1,5 +1,6 @@
 """Constants for the Adax integration."""
 
+import datetime
 from typing import Final
 
 ACCOUNT_ID: Final = "account_id"
@@ -9,3 +10,5 @@ DOMAIN: Final = "adax"
 LOCAL = "Local"
 WIFI_SSID = "wifi_ssid"
 WIFI_PSWD = "wifi_pswd"
+
+SCAN_INTERVAL = datetime.timedelta(seconds=60)

homeassistant/components/adax/coordinator.py

@@ -0,0 +1,71 @@
+"""DataUpdateCoordinator for the Adax component."""
+
+import logging
+from typing import Any, cast
+
+from adax import Adax
+from adax_local import Adax as AdaxLocal
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_IP_ADDRESS, CONF_PASSWORD, CONF_TOKEN
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import ACCOUNT_ID, SCAN_INTERVAL
+
+_LOGGER = logging.getLogger(__name__)
+
+
+type AdaxConfigEntry = ConfigEntry[AdaxCloudCoordinator | AdaxLocalCoordinator]
+
+
+class AdaxCloudCoordinator(DataUpdateCoordinator[dict[str, dict[str, Any]]]):
+    """Coordinator for updating data to and from Adax (cloud)."""
+
+    def __init__(self, hass: HomeAssistant, entry: AdaxConfigEntry) -> None:
+        """Initialize the Adax coordinator used for Cloud mode."""
+        super().__init__(
+            hass,
+            config_entry=entry,
+            logger=_LOGGER,
+            name="AdaxCloud",
+            update_interval=SCAN_INTERVAL,
+        )
+
+        self.adax_data_handler = Adax(
+            entry.data[ACCOUNT_ID],
+            entry.data[CONF_PASSWORD],
+            websession=async_get_clientsession(hass),
+        )
+
+    async def _async_update_data(self) -> dict[str, dict[str, Any]]:
+        """Fetch data from the Adax."""
+        rooms = await self.adax_data_handler.get_rooms() or []
+        return {r["id"]: r for r in rooms}
+
+
+class AdaxLocalCoordinator(DataUpdateCoordinator[dict[str, Any] | None]):
+    """Coordinator for updating data to and from Adax (local)."""
+
+    def __init__(self, hass: HomeAssistant, entry: AdaxConfigEntry) -> None:
+        """Initialize the Adax coordinator used for Local mode."""
+        super().__init__(
+            hass,
+            config_entry=entry,
+            logger=_LOGGER,
+            name="AdaxLocal",
+            update_interval=SCAN_INTERVAL,
+        )
+
+        self.adax_data_handler = AdaxLocal(
+            entry.data[CONF_IP_ADDRESS],
+            entry.data[CONF_TOKEN],
+            websession=async_get_clientsession(hass, verify_ssl=False),
+        )
+
+    async def _async_update_data(self) -> dict[str, Any]:
+        """Fetch data from the Adax."""
+        if result := await self.adax_data_handler.get_status():
+            return cast(dict[str, Any], result)
+        raise UpdateFailed("Got invalid status from device")

tests/components/adax/__init__.py

@@ -1 +1,12 @@
 """Tests for the Adax integration."""
+
+from homeassistant.core import HomeAssistant
+
+from tests.common import MockConfigEntry
+
+
+async def setup_integration(hass: HomeAssistant, entry: MockConfigEntry) -> None:
+    """Set up the Adax integration in Home Assistant."""
+    entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()

tests/components/adax/conftest.py

@@ -0,0 +1,89 @@
+"""Fixtures for Adax testing."""
+
+from typing import Any
+from unittest.mock import patch
+
+import pytest
+
+from homeassistant.components.adax.const import (
+    ACCOUNT_ID,
+    CLOUD,
+    CONNECTION_TYPE,
+    DOMAIN,
+    LOCAL,
+)
+from homeassistant.const import (
+    CONF_IP_ADDRESS,
+    CONF_PASSWORD,
+    CONF_TOKEN,
+    CONF_UNIQUE_ID,
+)
+
+from tests.common import AsyncMock, MockConfigEntry
+
+CLOUD_CONFIG = {
+    ACCOUNT_ID: 12345,
+    CONF_PASSWORD: "pswd",
+    CONNECTION_TYPE: CLOUD,
+}
+
+LOCAL_CONFIG = {
+    CONF_IP_ADDRESS: "192.168.1.12",
+    CONF_TOKEN: "TOKEN-123",
+    CONF_UNIQUE_ID: "11:22:33:44:55:66",
+    CONNECTION_TYPE: LOCAL,
+}
+
+
+CLOUD_DEVICE_DATA: dict[str, Any] = [
+    {
+        "id": "1",
+        "homeId": "1",
+        "name": "Room 1",
+        "temperature": 15,
+        "targetTemperature": 20,
+        "heatingEnabled": True,
+    }
+]
+
+LOCAL_DEVICE_DATA: dict[str, Any] = {
+    "current_temperature": 15,
+    "target_temperature": 20,
+}
+
+
+@pytest.fixture
+def mock_cloud_config_entry(request: pytest.FixtureRequest) -> MockConfigEntry:
+    """Mock a "CLOUD" config entry."""
+    return MockConfigEntry(domain=DOMAIN, data=CLOUD_CONFIG)
+
+
+@pytest.fixture
+def mock_local_config_entry(request: pytest.FixtureRequest) -> MockConfigEntry:
+    """Mock a "LOCAL" config entry."""
+    return MockConfigEntry(domain=DOMAIN, data=LOCAL_CONFIG)
+
+
+@pytest.fixture
+def mock_adax_cloud():
+    """Mock climate data."""
+    with patch("homeassistant.components.adax.coordinator.Adax") as mock_adax:
+        mock_adax_class = mock_adax.return_value
+
+        mock_adax_class.get_rooms = AsyncMock()
+        mock_adax_class.get_rooms.return_value = CLOUD_DEVICE_DATA
+
+        mock_adax_class.update = AsyncMock()
+        mock_adax_class.update.return_value = None
+        yield mock_adax_class
+
+
+@pytest.fixture
+def mock_adax_local():
+    """Mock climate data."""
+    with patch("homeassistant.components.adax.coordinator.AdaxLocal") as mock_adax:
+        mock_adax_class = mock_adax.return_value
+
+        mock_adax_class.get_status = AsyncMock()
+        mock_adax_class.get_status.return_value = LOCAL_DEVICE_DATA
+        yield mock_adax_class

tests/components/adax/test_climate.py

@@ -0,0 +1,85 @@
+"""Test Adax climate entity."""
+
+from homeassistant.components.adax.const import SCAN_INTERVAL
+from homeassistant.components.climate import ATTR_CURRENT_TEMPERATURE, HVACMode
+from homeassistant.const import ATTR_TEMPERATURE, STATE_UNAVAILABLE, Platform
+from homeassistant.core import HomeAssistant
+
+from . import setup_integration
+from .conftest import CLOUD_DEVICE_DATA, LOCAL_DEVICE_DATA
+
+from tests.common import AsyncMock, MockConfigEntry, async_fire_time_changed
+from tests.test_setup import FrozenDateTimeFactory
+
+
+async def test_climate_cloud(
+    hass: HomeAssistant,
+    freezer: FrozenDateTimeFactory,
+    mock_cloud_config_entry: MockConfigEntry,
+    mock_adax_cloud: AsyncMock,
+) -> None:
+    """Test states of the (cloud) Climate entity."""
+    await setup_integration(hass, mock_cloud_config_entry)
+    mock_adax_cloud.get_rooms.assert_called_once()
+
+    assert len(hass.states.async_entity_ids(Platform.CLIMATE)) == 1
+    entity_id = hass.states.async_entity_ids(Platform.CLIMATE)[0]
+
+    state = hass.states.get(entity_id)
+
+    assert state
+    assert state.state == HVACMode.HEAT
+    assert (
+        state.attributes[ATTR_TEMPERATURE] == CLOUD_DEVICE_DATA[0]["targetTemperature"]
+    )
+    assert (
+        state.attributes[ATTR_CURRENT_TEMPERATURE]
+        == CLOUD_DEVICE_DATA[0]["temperature"]
+    )
+
+    mock_adax_cloud.get_rooms.side_effect = Exception()
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == STATE_UNAVAILABLE
+
+
+async def test_climate_local(
+    hass: HomeAssistant,
+    freezer: FrozenDateTimeFactory,
+    mock_local_config_entry: MockConfigEntry,
+    mock_adax_local: AsyncMock,
+) -> None:
+    """Test states of the (local) Climate entity."""
+    await setup_integration(hass, mock_local_config_entry)
+    mock_adax_local.get_status.assert_called_once()
+
+    assert len(hass.states.async_entity_ids(Platform.CLIMATE)) == 1
+    entity_id = hass.states.async_entity_ids(Platform.CLIMATE)[0]
+
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == HVACMode.HEAT
+    assert (
+        state.attributes[ATTR_TEMPERATURE] == (LOCAL_DEVICE_DATA["target_temperature"])
+    )
+    assert (
+        state.attributes[ATTR_CURRENT_TEMPERATURE]
+        == (LOCAL_DEVICE_DATA["current_temperature"])
+    )
+
+    mock_adax_local.get_status.side_effect = Exception()
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == STATE_UNAVAILABLE