Move atag coordinator to separate class (#127071)

homeassistant/components/atag/__init__.py

@@ -1,18 +1,10 @@
 """The ATAG Integration."""
 
-from asyncio import timeout
-from datetime import timedelta
-import logging
-
-from pyatag import AtagException, AtagOne
-
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-_LOGGER = logging.getLogger(__name__)
+from .coordinator import AtagDataUpdateCoordinator
 
 DOMAIN = "atag"
 PLATFORMS = [Platform.CLIMATE, Platform.SENSOR, Platform.WATER_HEATER]
@@ -21,31 +13,12 @@ PLATFORMS = [Platform.CLIMATE, Platform.SENSOR, Platform.WATER_HEATER]
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Atag integration from a config entry."""
 
-    async def _async_update_data():
+    coordinator = AtagDataUpdateCoordinator(hass, entry)
-        """Update data via library."""
-        async with timeout(20):
-            try:
-                await atag.update()
-            except AtagException as err:
-                raise UpdateFailed(err) from err
-        return atag
-
-    atag = AtagOne(
-        session=async_get_clientsession(hass), **entry.data, device=entry.unique_id
-    )
-    coordinator = DataUpdateCoordinator[AtagOne](
-        hass,
-        _LOGGER,
-        name=DOMAIN.title(),
-        update_method=_async_update_data,
-        update_interval=timedelta(seconds=60),
-    )
-
     await coordinator.async_config_entry_first_refresh()
 
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
     if entry.unique_id is None:
-        hass.config_entries.async_update_entry(entry, unique_id=atag.id)
+        hass.config_entries.async_update_entry(entry, unique_id=coordinator.atag.id)
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 

homeassistant/components/atag/climate.py

@@ -53,46 +53,46 @@ class AtagThermostat(AtagEntity, ClimateEntity):
     def __init__(self, coordinator, atag_id):
         """Initialize an Atag climate device."""
         super().__init__(coordinator, atag_id)
-        self._attr_temperature_unit = coordinator.data.climate.temp_unit
+        self._attr_temperature_unit = coordinator.atag.climate.temp_unit
 
     @property
     def hvac_mode(self) -> HVACMode | None:
         """Return hvac operation ie. heat, cool mode."""
-        return try_parse_enum(HVACMode, self.coordinator.data.climate.hvac_mode)
+        return try_parse_enum(HVACMode, self.coordinator.atag.climate.hvac_mode)
 
     @property
     def hvac_action(self) -> HVACAction | None:
         """Return the current running hvac operation."""
-        is_active = self.coordinator.data.climate.status
+        is_active = self.coordinator.atag.climate.status
         return HVACAction.HEATING if is_active else HVACAction.IDLE
 
     @property
     def current_temperature(self) -> float | None:
         """Return the current temperature."""
-        return self.coordinator.data.climate.temperature
+        return self.coordinator.atag.climate.temperature
 
     @property
     def target_temperature(self) -> float | None:
         """Return the temperature we try to reach."""
-        return self.coordinator.data.climate.target_temperature
+        return self.coordinator.atag.climate.target_temperature
 
     @property
     def preset_mode(self) -> str | None:
         """Return the current preset mode, e.g., auto, manual, fireplace, extend, etc."""
-        preset = self.coordinator.data.climate.preset_mode
+        preset = self.coordinator.atag.climate.preset_mode
         return PRESET_INVERTED.get(preset)
 
     async def async_set_temperature(self, **kwargs: Any) -> None:
         """Set new target temperature."""
-        await self.coordinator.data.climate.set_temp(kwargs.get(ATTR_TEMPERATURE))
+        await self.coordinator.atag.climate.set_temp(kwargs.get(ATTR_TEMPERATURE))
         self.async_write_ha_state()
 
     async def async_set_hvac_mode(self, hvac_mode: HVACMode) -> None:
         """Set new target hvac mode."""
-        await self.coordinator.data.climate.set_hvac_mode(hvac_mode)
+        await self.coordinator.atag.climate.set_hvac_mode(hvac_mode)
         self.async_write_ha_state()
 
     async def async_set_preset_mode(self, preset_mode: str) -> None:
         """Set new preset mode."""
-        await self.coordinator.data.climate.set_preset_mode(PRESET_MAP[preset_mode])
+        await self.coordinator.atag.climate.set_preset_mode(PRESET_MAP[preset_mode])
         self.async_write_ha_state()

homeassistant/components/atag/coordinator.py

@@ -0,0 +1,39 @@
+"""The ATAG Integration."""
+
+from asyncio import timeout
+from datetime import timedelta
+import logging
+
+from pyatag import AtagException, AtagOne
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class AtagDataUpdateCoordinator(DataUpdateCoordinator[None]):
+    """Atag data update coordinator."""
+
+    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
+        """Initialize Atag coordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            name="Atag",
+            update_interval=timedelta(seconds=60),
+        )
+
+        self.atag = AtagOne(
+            session=async_get_clientsession(hass), **entry.data, device=entry.unique_id
+        )
+
+    async def _async_update_data(self) -> None:
+        """Update data via library."""
+        async with timeout(20):
+            try:
+                await self.atag.update()
+            except AtagException as err:
+                raise UpdateFailed(err) from err

homeassistant/components/atag/entity.py

@@ -1,36 +1,30 @@
 """The ATAG Integration."""
 
-from pyatag import AtagOne
-
 from homeassistant.helpers.device_registry import DeviceInfo
-from homeassistant.helpers.update_coordinator import (
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-)
 
 from . import DOMAIN
+from .coordinator import AtagDataUpdateCoordinator
 
 
-class AtagEntity(CoordinatorEntity[DataUpdateCoordinator[AtagOne]]):
+class AtagEntity(CoordinatorEntity[AtagDataUpdateCoordinator]):
     """Defines a base Atag entity."""
 
-    def __init__(
+    def __init__(self, coordinator: AtagDataUpdateCoordinator, atag_id: str) -> None:
-        self, coordinator: DataUpdateCoordinator[AtagOne], atag_id: str
-    ) -> None:
         """Initialize the Atag entity."""
         super().__init__(coordinator)
 
         self._id = atag_id
         self._attr_name = DOMAIN.title()
-        self._attr_unique_id = f"{coordinator.data.id}-{atag_id}"
+        self._attr_unique_id = f"{coordinator.atag.id}-{atag_id}"
 
     @property
     def device_info(self) -> DeviceInfo:
         """Return info for device registry."""
         return DeviceInfo(
-            identifiers={(DOMAIN, self.coordinator.data.id)},
+            identifiers={(DOMAIN, self.coordinator.atag.id)},
             manufacturer="Atag",
             model="Atag One",
             name="Atag Thermostat",
-            sw_version=self.coordinator.data.apiversion,
+            sw_version=self.coordinator.atag.apiversion,
         )

homeassistant/components/atag/sensor.py

@@ -43,28 +43,28 @@ class AtagSensor(AtagEntity, SensorEntity):
         """Initialize Atag sensor."""
         super().__init__(coordinator, SENSORS[sensor])
         self._attr_name = sensor
-        if coordinator.data.report[self._id].sensorclass in (
+        if coordinator.atag.report[self._id].sensorclass in (
             SensorDeviceClass.PRESSURE,
             SensorDeviceClass.TEMPERATURE,
         ):
-            self._attr_device_class = coordinator.data.report[self._id].sensorclass
+            self._attr_device_class = coordinator.atag.report[self._id].sensorclass
-        if coordinator.data.report[self._id].measure in (
+        if coordinator.atag.report[self._id].measure in (
             UnitOfPressure.BAR,
             UnitOfTemperature.CELSIUS,
             UnitOfTemperature.FAHRENHEIT,
             PERCENTAGE,
             UnitOfTime.HOURS,
         ):
-            self._attr_native_unit_of_measurement = coordinator.data.report[
+            self._attr_native_unit_of_measurement = coordinator.atag.report[
                 self._id
             ].measure
 
     @property
     def native_value(self):
         """Return the state of the sensor."""
-        return self.coordinator.data.report[self._id].state
+        return self.coordinator.atag.report[self._id].state
 
     @property
     def icon(self):
         """Return icon."""
-        return self.coordinator.data.report[self._id].icon
+        return self.coordinator.atag.report[self._id].icon

homeassistant/components/atag/water_heater.py

@@ -37,30 +37,30 @@ class AtagWaterHeater(AtagEntity, WaterHeaterEntity):
     @property
     def current_temperature(self):
         """Return the current temperature."""
-        return self.coordinator.data.dhw.temperature
+        return self.coordinator.atag.dhw.temperature
 
     @property
     def current_operation(self):
         """Return current operation."""
-        operation = self.coordinator.data.dhw.current_operation
+        operation = self.coordinator.atag.dhw.current_operation
         return operation if operation in self.operation_list else STATE_OFF
 
     async def async_set_temperature(self, **kwargs: Any) -> None:
         """Set new target temperature."""
-        if await self.coordinator.data.dhw.set_temp(kwargs.get(ATTR_TEMPERATURE)):
+        if await self.coordinator.atag.dhw.set_temp(kwargs.get(ATTR_TEMPERATURE)):
             self.async_write_ha_state()
 
     @property
     def target_temperature(self):
         """Return the setpoint if water demand, otherwise return base temp (comfort level)."""
-        return self.coordinator.data.dhw.target_temperature
+        return self.coordinator.atag.dhw.target_temperature
 
     @property
     def max_temp(self) -> float:
         """Return the maximum temperature."""
-        return self.coordinator.data.dhw.max_temp
+        return self.coordinator.atag.dhw.max_temp
 
     @property
     def min_temp(self) -> float:
         """Return the minimum temperature."""
-        return self.coordinator.data.dhw.min_temp
+        return self.coordinator.atag.dhw.min_temp

tests/components/atag/__init__.py

@@ -1,6 +1,8 @@
 """Tests for the Atag integration."""
 
-from homeassistant.components.atag import DOMAIN, AtagException
+from pyatag import AtagException
+
+from homeassistant.components.atag import DOMAIN
 from homeassistant.const import CONF_HOST, CONF_PORT
 from homeassistant.core import HomeAssistant
 

tests/components/atag/test_climate.py

@@ -110,4 +110,4 @@ async def test_update_failed(
         await hass.async_block_till_done()
         updater.assert_called_once()
         assert not coordinator.last_update_success
-        assert coordinator.data.id == UID
+        assert coordinator.atag.id == UID