Huum - Introduce coordinator to support multiple platforms (home-assistant#148889)

Co-authored-by: Josef Zweck <[email protected]>

Author: vincentwolsink

CODEOWNERS

@@ -2,46 +2,28 @@
 
 from __future__ import annotations
 
-import logging
-
-from huum.exceptions import Forbidden, NotAuthenticated
-from huum.huum import Huum
-
-from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
-
-from .const import DOMAIN, PLATFORMS
 
-_LOGGER = logging.getLogger(__name__)
+from .const import PLATFORMS
+from .coordinator import HuumConfigEntry, HuumDataUpdateCoordinator
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, config_entry: HuumConfigEntry) -> bool:
     """Set up Huum from a config entry."""
-    username = entry.data[CONF_USERNAME]
-    password = entry.data[CONF_PASSWORD]
+    coordinator = HuumDataUpdateCoordinator(
+        hass=hass,
+        config_entry=config_entry,
+    )
 
-    huum = Huum(username, password, session=async_get_clientsession(hass))
+    await coordinator.async_config_entry_first_refresh()
+    config_entry.runtime_data = coordinator
 
-    try:
-        await huum.status()
-    except (Forbidden, NotAuthenticated) as err:
-        _LOGGER.error("Could not log in to Huum with given credentials")
-        raise ConfigEntryNotReady(
-            "Could not log in to Huum with given credentials"
-        ) from err
-
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = huum
-
-    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+    await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(
+    hass: HomeAssistant, config_entry: HuumConfigEntry
+) -> bool:
     """Unload a config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS)

homeassistant/components/huum/__init__.py

@@ -7,38 +7,35 @@
 
 from huum.const import SaunaStatus
 from huum.exceptions import SafetyException
-from huum.huum import Huum
-from huum.schemas import HuumStatusResponse
 
 from homeassistant.components.climate import (
     ClimateEntity,
     ClimateEntityFeature,
     HVACMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TEMPERATURE, PRECISION_WHOLE, UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import DOMAIN
+from .coordinator import HuumConfigEntry, HuumDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: HuumConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the Huum sauna with config flow."""
-    huum_handler = hass.data.setdefault(DOMAIN, {})[entry.entry_id]
+    async_add_entities([HuumDevice(entry.runtime_data)])
 
-    async_add_entities([HuumDevice(huum_handler, entry.entry_id)], True)
 
-
-class HuumDevice(ClimateEntity):
+class HuumDevice(CoordinatorEntity[HuumDataUpdateCoordinator], ClimateEntity):
     """Representation of a heater."""
 
     _attr_hvac_modes = [HVACMode.HEAT, HVACMode.OFF]
@@ -54,24 +51,22 @@ class HuumDevice(ClimateEntity):
     _attr_has_entity_name = True
     _attr_name = None
 
-    _target_temperature: int | None = None
-    _status: HuumStatusResponse | None = None
-
-    def __init__(self, huum_handler: Huum, unique_id: str) -> None:
+    def __init__(self, coordinator: HuumDataUpdateCoordinator) -> None:
         """Initialize the heater."""
-        self._attr_unique_id = unique_id
+        super().__init__(coordinator)
+
+        self._attr_unique_id = coordinator.config_entry.entry_id
         self._attr_device_info = DeviceInfo(
-            identifiers={(DOMAIN, unique_id)},
+            identifiers={(DOMAIN, coordinator.config_entry.entry_id)},
             name="Huum sauna",
             manufacturer="Huum",
+            model="UKU WiFi",
         )
 
-        self._huum_handler = huum_handler
-
     @property
     def hvac_mode(self) -> HVACMode:
         """Return hvac operation ie. heat, cool mode."""
-        if self._status and self._status.status == SaunaStatus.ONLINE_HEATING:
+        if self.coordinator.data.status == SaunaStatus.ONLINE_HEATING:
             return HVACMode.HEAT
         return HVACMode.OFF
 
@@ -85,41 +80,33 @@ def icon(self) -> str:
     @property
     def current_temperature(self) -> int | None:
         """Return the current temperature."""
-        if (status := self._status) is not None:
-            return status.temperature
-        return None
+        return self.coordinator.data.temperature
 
     @property
     def target_temperature(self) -> int:
         """Return the temperature we try to reach."""
-        return self._target_temperature or int(self.min_temp)
+        return self.coordinator.data.target_temperature or int(self.min_temp)
 
     async def async_set_hvac_mode(self, hvac_mode: HVACMode) -> None:
         """Set hvac mode."""
         if hvac_mode == HVACMode.HEAT:
             await self._turn_on(self.target_temperature)
         elif hvac_mode == HVACMode.OFF:
-            await self._huum_handler.turn_off()
+            await self.coordinator.huum.turn_off()
+        await self.coordinator.async_refresh()
 
     async def async_set_temperature(self, **kwargs: Any) -> None:
         """Set new target temperature."""
         temperature = kwargs.get(ATTR_TEMPERATURE)
-        if temperature is None:
+        if temperature is None or self.hvac_mode != HVACMode.HEAT:
             return
-        self._target_temperature = temperature
-
-        if self.hvac_mode == HVACMode.HEAT:
-            await self._turn_on(temperature)
 
-    async def async_update(self) -> None:
-        """Get the latest status data."""
-        self._status = await self._huum_handler.status()
-        if self._target_temperature is None or self.hvac_mode == HVACMode.HEAT:
-            self._target_temperature = self._status.target_temperature
+        await self._turn_on(temperature)
+        await self.coordinator.async_refresh()
 
     async def _turn_on(self, temperature: int) -> None:
         try:
-            await self._huum_handler.turn_on(temperature)
+            await self.coordinator.huum.turn_on(temperature)
         except (ValueError, SafetyException) as err:
             _LOGGER.error(str(err))
             raise HomeAssistantError(f"Unable to turn on sauna: {err}") from err

homeassistant/components/huum/climate.py

@@ -37,12 +37,12 @@ async def async_step_user(
         errors = {}
         if user_input is not None:
             try:
-                huum_handler = Huum(
+                huum = Huum(
                     user_input[CONF_USERNAME],
                     user_input[CONF_PASSWORD],
                     session=async_get_clientsession(self.hass),
                 )
-                await huum_handler.status()
+                await huum.status()
             except (Forbidden, NotAuthenticated):
                 # Most likely Forbidden as that is what is returned from `.status()` with bad creds
                 _LOGGER.error("Could not log in to Huum with given credentials")

homeassistant/components/huum/config_flow.py

@@ -0,0 +1,60 @@
+"""DataUpdateCoordinator for Huum."""
+
+from __future__ import annotations
+
+from datetime import timedelta
+import logging
+
+from huum.exceptions import Forbidden, NotAuthenticated
+from huum.huum import Huum
+from huum.schemas import HuumStatusResponse
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
+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 DOMAIN
+
+type HuumConfigEntry = ConfigEntry[HuumDataUpdateCoordinator]
+
+_LOGGER = logging.getLogger(__name__)
+UPDATE_INTERVAL = timedelta(seconds=30)
+
+
+class HuumDataUpdateCoordinator(DataUpdateCoordinator[HuumStatusResponse]):
+    """Class to manage fetching data from the API."""
+
+    config_entry: HuumConfigEntry
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        config_entry: HuumConfigEntry,
+    ) -> None:
+        """Initialize."""
+        super().__init__(
+            hass=hass,
+            logger=_LOGGER,
+            name=DOMAIN,
+            update_interval=UPDATE_INTERVAL,
+            config_entry=config_entry,
+        )
+
+        self.huum = Huum(
+            config_entry.data[CONF_USERNAME],
+            config_entry.data[CONF_PASSWORD],
+            session=async_get_clientsession(hass),
+        )
+
+    async def _async_update_data(self) -> HuumStatusResponse:
+        """Get the latest status data."""
+
+        try:
+            return await self.huum.status()
+        except (Forbidden, NotAuthenticated) as err:
+            _LOGGER.error("Could not log in to Huum with given credentials")
+            raise UpdateFailed(
+                "Could not log in to Huum with given credentials"
+            ) from err

homeassistant/components/huum/coordinator.py

@@ -1,7 +1,7 @@
 {
   "domain": "huum",
   "name": "Huum",
-  "codeowners": ["@frwickst"],
+  "codeowners": ["@frwickst", "@vincentwolsink"],
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/huum",
   "iot_class": "cloud_polling",

homeassistant/components/huum/manifest.json

@@ -1 +1,18 @@
 """Tests for the huum integration."""
+
+from unittest.mock import patch
+
+from homeassistant.const import Platform
+from homeassistant.core import HomeAssistant
+
+from tests.common import MockConfigEntry
+
+
+async def setup_with_selected_platforms(
+    hass: HomeAssistant, entry: MockConfigEntry, platforms: list[Platform]
+) -> None:
+    """Set up the Huum integration with the selected platforms."""
+    entry.add_to_hass(hass)
+    with patch("homeassistant.components.huum.PLATFORMS", platforms):
+        assert await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()

tests/components/huum/__init__.py

@@ -0,0 +1,72 @@
+"""Configuration for Huum tests."""
+
+from collections.abc import Generator
+from unittest.mock import AsyncMock, patch
+
+from huum.const import SaunaStatus
+import pytest
+
+from homeassistant.components.huum.const import DOMAIN
+from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
+
+from tests.common import MockConfigEntry
+
+
+@pytest.fixture
+def mock_huum() -> Generator[AsyncMock]:
+    """Mock data from the API."""
+    huum = AsyncMock()
+    with (
+        patch(
+            "homeassistant.components.huum.config_flow.Huum.status",
+            return_value=huum,
+        ),
+        patch(
+            "homeassistant.components.huum.coordinator.Huum.status",
+            return_value=huum,
+        ),
+        patch(
+            "homeassistant.components.huum.coordinator.Huum.turn_on",
+            return_value=huum,
+        ) as turn_on,
+    ):
+        huum.status = SaunaStatus.ONLINE_NOT_HEATING
+        huum.door_closed = True
+        huum.temperature = 30
+        huum.sauna_name = 123456
+        huum.target_temperature = 80
+        huum.light = 1
+        huum.humidity = 5
+        huum.sauna_config.child_lock = "OFF"
+        huum.sauna_config.max_heating_time = 3
+        huum.sauna_config.min_heating_time = 0
+        huum.sauna_config.max_temp = 110
+        huum.sauna_config.min_temp = 40
+        huum.sauna_config.max_timer = 0
+        huum.sauna_config.min_timer = 0
+        huum.turn_on = turn_on
+
+        yield huum
+
+
+@pytest.fixture
+def mock_setup_entry() -> Generator[AsyncMock]:
+    """Mock setting up a config entry."""
+    with patch(
+        "homeassistant.components.huum.async_setup_entry", return_value=True
+    ) as setup_entry_mock:
+        yield setup_entry_mock
+
+
+@pytest.fixture
+def mock_config_entry() -> MockConfigEntry:
+    """Mock a config entry."""
+    return MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_USERNAME: "[email protected]",
+            CONF_PASSWORD: "ukuuku",
+        },
+        unique_id="123456",
+        entry_id="AABBCC112233",
+    )