Merge coordinators in Airgradient (#124714)

homeassistant/components/airgradient/__init__.py

@@ -2,18 +2,14 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
-
-from airgradient import AirGradientClient, get_model_name
+from airgradient import AirGradientClient
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import DOMAIN
-from .coordinator import AirGradientConfigCoordinator, AirGradientMeasurementCoordinator
+from .coordinator import AirGradientCoordinator
 
 PLATFORMS: list[Platform] = [
     Platform.BUTTON,
@@ -25,15 +21,7 @@ PLATFORMS: list[Platform] = [
 ]
 
 
-@dataclass
-class AirGradientData:
-    """AirGradient data class."""
-
-    measurement: AirGradientMeasurementCoordinator
-    config: AirGradientConfigCoordinator
-
-
-type AirGradientConfigEntry = ConfigEntry[AirGradientData]
+type AirGradientConfigEntry = ConfigEntry[AirGradientCoordinator]
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: AirGradientConfigEntry) -> bool:
@@ -43,27 +31,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: AirGradientConfigEntry)
         entry.data[CONF_HOST], session=async_get_clientsession(hass)
     )
 
-    measurement_coordinator = AirGradientMeasurementCoordinator(hass, client)
-    config_coordinator = AirGradientConfigCoordinator(hass, client)
+    coordinator = AirGradientCoordinator(hass, client)
 
-    await measurement_coordinator.async_config_entry_first_refresh()
-    await config_coordinator.async_config_entry_first_refresh()
+    await coordinator.async_config_entry_first_refresh()
 
-    device_registry = dr.async_get(hass)
-    device_registry.async_get_or_create(
-        config_entry_id=entry.entry_id,
-        identifiers={(DOMAIN, measurement_coordinator.serial_number)},
-        manufacturer="AirGradient",
-        model=get_model_name(measurement_coordinator.data.model),
-        model_id=measurement_coordinator.data.model,
-        serial_number=measurement_coordinator.data.serial_number,
-        sw_version=measurement_coordinator.data.firmware_version,
-    )
-
-    entry.runtime_data = AirGradientData(
-        measurement=measurement_coordinator,
-        config=config_coordinator,
-    )
+    entry.runtime_data = coordinator
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 

homeassistant/components/airgradient/button.py

@@ -15,8 +15,9 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import DOMAIN, AirGradientConfigEntry
-from .coordinator import AirGradientConfigCoordinator
+from . import AirGradientConfigEntry
+from .const import DOMAIN
+from .coordinator import AirGradientCoordinator
 from .entity import AirGradientEntity
 
 
@@ -47,8 +48,8 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up AirGradient button entities based on a config entry."""
-    model = entry.runtime_data.measurement.data.model
-    coordinator = entry.runtime_data.config
+    coordinator = entry.runtime_data
+    model = coordinator.data.measures.model
 
     added_entities = False
 
@@ -57,7 +58,7 @@ async def async_setup_entry(
         nonlocal added_entities
 
         if (
-            coordinator.data.configuration_control is ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control is ConfigurationControl.LOCAL
             and not added_entities
         ):
             entities = [AirGradientButton(coordinator, CO2_CALIBRATION)]
@@ -67,7 +68,8 @@ async def async_setup_entry(
             async_add_entities(entities)
             added_entities = True
         elif (
-            coordinator.data.configuration_control is not ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control
+            is not ConfigurationControl.LOCAL
             and added_entities
         ):
             entity_registry = er.async_get(hass)
@@ -87,11 +89,10 @@ class AirGradientButton(AirGradientEntity, ButtonEntity):
     """Defines an AirGradient button."""
 
     entity_description: AirGradientButtonEntityDescription
-    coordinator: AirGradientConfigCoordinator
 
     def __init__(
         self,
-        coordinator: AirGradientConfigCoordinator,
+        coordinator: AirGradientCoordinator,
         description: AirGradientButtonEntityDescription,
     ) -> None:
         """Initialize airgradient button."""

homeassistant/components/airgradient/coordinator.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from datetime import timedelta
 from typing import TYPE_CHECKING
 
@@ -16,7 +17,15 @@ if TYPE_CHECKING:
     from . import AirGradientConfigEntry
 
 
-class AirGradientCoordinator[_DataT](DataUpdateCoordinator[_DataT]):
+@dataclass
+class AirGradientData:
+    """Class for AirGradient data."""
+
+    measures: Measures
+    config: Config
+
+
+class AirGradientCoordinator(DataUpdateCoordinator[AirGradientData]):
     """Class to manage fetching AirGradient data."""
 
     config_entry: AirGradientConfigEntry
@@ -33,25 +42,11 @@ class AirGradientCoordinator[_DataT](DataUpdateCoordinator[_DataT]):
         assert self.config_entry.unique_id
         self.serial_number = self.config_entry.unique_id
 
-    async def _async_update_data(self) -> _DataT:
+    async def _async_update_data(self) -> AirGradientData:
         try:
-            return await self._update_data()
+            measures = await self.client.get_current_measures()
+            config = await self.client.get_config()
         except AirGradientError as error:
             raise UpdateFailed(error) from error
-
-    async def _update_data(self) -> _DataT:
-        raise NotImplementedError
-
-
-class AirGradientMeasurementCoordinator(AirGradientCoordinator[Measures]):
-    """Class to manage fetching AirGradient data."""
-
-    async def _update_data(self) -> Measures:
-        return await self.client.get_current_measures()
-
-
-class AirGradientConfigCoordinator(AirGradientCoordinator[Config]):
-    """Class to manage fetching AirGradient data."""
-
-    async def _update_data(self) -> Config:
-        return await self.client.get_config()
+        else:
+            return AirGradientData(measures, config)

homeassistant/components/airgradient/entity.py

@@ -1,5 +1,7 @@
 """Base class for AirGradient entities."""
 
+from airgradient import get_model_name
+
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
@@ -15,6 +17,12 @@ class AirGradientEntity(CoordinatorEntity[AirGradientCoordinator]):
     def __init__(self, coordinator: AirGradientCoordinator) -> None:
         """Initialize airgradient entity."""
         super().__init__(coordinator)
+        measures = coordinator.data.measures
         self._attr_device_info = DeviceInfo(
             identifiers={(DOMAIN, coordinator.serial_number)},
+            manufacturer="AirGradient",
+            model=get_model_name(measures.model),
+            model_id=measures.model,
+            serial_number=coordinator.serial_number,
+            sw_version=measures.firmware_version,
         )

homeassistant/components/airgradient/number.py

@@ -18,7 +18,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import AirGradientConfigEntry
 from .const import DOMAIN
-from .coordinator import AirGradientConfigCoordinator
+from .coordinator import AirGradientCoordinator
 from .entity import AirGradientEntity
 
 
@@ -62,8 +62,8 @@ async def async_setup_entry(
 ) -> None:
     """Set up AirGradient number entities based on a config entry."""
 
-    model = entry.runtime_data.measurement.data.model
-    coordinator = entry.runtime_data.config
+    coordinator = entry.runtime_data
+    model = coordinator.data.measures.model
 
     added_entities = False
 
@@ -72,7 +72,7 @@ async def async_setup_entry(
         nonlocal added_entities
 
         if (
-            coordinator.data.configuration_control is ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control is ConfigurationControl.LOCAL
             and not added_entities
         ):
             entities = []
@@ -84,7 +84,8 @@ async def async_setup_entry(
             async_add_entities(entities)
             added_entities = True
         elif (
-            coordinator.data.configuration_control is not ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control
+            is not ConfigurationControl.LOCAL
             and added_entities
         ):
             entity_registry = er.async_get(hass)
@@ -104,11 +105,10 @@ class AirGradientNumber(AirGradientEntity, NumberEntity):
     """Defines an AirGradient number entity."""
 
     entity_description: AirGradientNumberEntityDescription
-    coordinator: AirGradientConfigCoordinator
 
     def __init__(
         self,
-        coordinator: AirGradientConfigCoordinator,
+        coordinator: AirGradientCoordinator,
         description: AirGradientNumberEntityDescription,
     ) -> None:
         """Initialize AirGradient number."""
@@ -119,7 +119,7 @@ class AirGradientNumber(AirGradientEntity, NumberEntity):
     @property
     def native_value(self) -> int | None:
         """Return the state of the number."""
-        return self.entity_description.value_fn(self.coordinator.data)
+        return self.entity_description.value_fn(self.coordinator.data.config)
 
     async def async_set_native_value(self, value: float) -> None:
         """Set the selected value."""

homeassistant/components/airgradient/select.py

@@ -18,7 +18,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import AirGradientConfigEntry
 from .const import DOMAIN, PM_STANDARD, PM_STANDARD_REVERSE
-from .coordinator import AirGradientConfigCoordinator
+from .coordinator import AirGradientCoordinator
 from .entity import AirGradientEntity
 
 
@@ -144,13 +144,11 @@ async def async_setup_entry(
 ) -> None:
     """Set up AirGradient select entities based on a config entry."""
 
-    coordinator = entry.runtime_data.config
-    measurement_coordinator = entry.runtime_data.measurement
+    coordinator = entry.runtime_data
+    model = coordinator.data.measures.model
 
     async_add_entities([AirGradientSelect(coordinator, CONFIG_CONTROL_ENTITY)])
 
-    model = measurement_coordinator.data.model
-
     added_entities = False
 
     @callback
@@ -158,7 +156,7 @@ async def async_setup_entry(
         nonlocal added_entities
 
         if (
-            coordinator.data.configuration_control is ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control is ConfigurationControl.LOCAL
             and not added_entities
         ):
             entities: list[AirGradientSelect] = [
@@ -179,7 +177,8 @@ async def async_setup_entry(
             async_add_entities(entities)
             added_entities = True
         elif (
-            coordinator.data.configuration_control is not ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control
+            is not ConfigurationControl.LOCAL
             and added_entities
         ):
             entity_registry = er.async_get(hass)
@@ -201,11 +200,10 @@ class AirGradientSelect(AirGradientEntity, SelectEntity):
     """Defines an AirGradient select entity."""
 
     entity_description: AirGradientSelectEntityDescription
-    coordinator: AirGradientConfigCoordinator
 
     def __init__(
         self,
-        coordinator: AirGradientConfigCoordinator,
+        coordinator: AirGradientCoordinator,
         description: AirGradientSelectEntityDescription,
     ) -> None:
         """Initialize AirGradient select."""
@@ -216,7 +214,7 @@ class AirGradientSelect(AirGradientEntity, SelectEntity):
     @property
     def current_option(self) -> str | None:
         """Return the state of the select."""
-        return self.entity_description.value_fn(self.coordinator.data)
+        return self.entity_description.value_fn(self.coordinator.data.config)
 
     async def async_select_option(self, option: str) -> None:
         """Change the selected option."""

homeassistant/components/airgradient/sensor.py

@@ -32,7 +32,7 @@ from homeassistant.helpers.typing import StateType
 
 from . import AirGradientConfigEntry
 from .const import PM_STANDARD, PM_STANDARD_REVERSE
-from .coordinator import AirGradientConfigCoordinator, AirGradientMeasurementCoordinator
+from .coordinator import AirGradientCoordinator
 from .entity import AirGradientEntity
 
 
@@ -218,7 +218,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up AirGradient sensor entities based on a config entry."""
 
-    coordinator = entry.runtime_data.measurement
+    coordinator = entry.runtime_data
     listener: Callable[[], None] | None = None
     not_setup: set[AirGradientMeasurementSensorEntityDescription] = set(
         MEASUREMENT_SENSOR_TYPES
@@ -232,7 +232,7 @@ async def async_setup_entry(
         not_setup = set()
         sensors = []
         for description in sensor_descriptions:
-            if description.value_fn(coordinator.data) is None:
+            if description.value_fn(coordinator.data.measures) is None:
                 not_setup.add(description)
             else:
                 sensors.append(AirGradientMeasurementSensor(coordinator, description))
@@ -248,64 +248,65 @@ async def async_setup_entry(
     add_entities()
 
     entities = [
-        AirGradientConfigSensor(entry.runtime_data.config, description)
+        AirGradientConfigSensor(coordinator, description)
         for description in CONFIG_SENSOR_TYPES
     ]
-    if "L" in coordinator.data.model:
+    if "L" in coordinator.data.measures.model:
         entities.extend(
-            AirGradientConfigSensor(entry.runtime_data.config, description)
+            AirGradientConfigSensor(coordinator, description)
             for description in CONFIG_LED_BAR_SENSOR_TYPES
         )
-    if "I" in coordinator.data.model:
+    if "I" in coordinator.data.measures.model:
         entities.extend(
-            AirGradientConfigSensor(entry.runtime_data.config, description)
+            AirGradientConfigSensor(coordinator, description)
             for description in CONFIG_DISPLAY_SENSOR_TYPES
         )
     async_add_entities(entities)
 
 
-class AirGradientMeasurementSensor(AirGradientEntity, SensorEntity):
+class AirGradientSensor(AirGradientEntity, SensorEntity):
     """Defines an AirGradient sensor."""
 
-    entity_description: AirGradientMeasurementSensorEntityDescription
-    coordinator: AirGradientMeasurementCoordinator
-
     def __init__(
         self,
-        coordinator: AirGradientMeasurementCoordinator,
-        description: AirGradientMeasurementSensorEntityDescription,
+        coordinator: AirGradientCoordinator,
+        description: SensorEntityDescription,
     ) -> None:
         """Initialize airgradient sensor."""
         super().__init__(coordinator)
         self.entity_description = description
         self._attr_unique_id = f"{coordinator.serial_number}-{description.key}"
 
+
+class AirGradientMeasurementSensor(AirGradientSensor):
+    """Defines an AirGradient sensor."""
+
+    entity_description: AirGradientMeasurementSensorEntityDescription
+
     @property
     def native_value(self) -> StateType:
         """Return the state of the sensor."""
-        return self.entity_description.value_fn(self.coordinator.data)
+        return self.entity_description.value_fn(self.coordinator.data.measures)
 
 
-class AirGradientConfigSensor(AirGradientEntity, SensorEntity):
+class AirGradientConfigSensor(AirGradientSensor):
     """Defines an AirGradient sensor."""
 
     entity_description: AirGradientConfigSensorEntityDescription
-    coordinator: AirGradientConfigCoordinator
 
     def __init__(
         self,
-        coordinator: AirGradientConfigCoordinator,
+        coordinator: AirGradientCoordinator,
         description: AirGradientConfigSensorEntityDescription,
     ) -> None:
         """Initialize airgradient sensor."""
-        super().__init__(coordinator)
-        self.entity_description = description
-        self._attr_unique_id = f"{coordinator.serial_number}-{description.key}"
+        super().__init__(coordinator, description)
         self._attr_entity_registry_enabled_default = (
-            coordinator.data.configuration_control is not ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control
+            is not ConfigurationControl.LOCAL
         )
 
     @property
     def native_value(self) -> StateType:
         """Return the state of the sensor."""
-        return self.entity_description.value_fn(self.coordinator.data)
+        return self.entity_description.value_fn(self.coordinator.data.config)

homeassistant/components/airgradient/switch.py

@@ -19,7 +19,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import AirGradientConfigEntry
 from .const import DOMAIN
-from .coordinator import AirGradientConfigCoordinator
+from .coordinator import AirGradientCoordinator
 from .entity import AirGradientEntity
 
 
@@ -46,7 +46,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up AirGradient switch entities based on a config entry."""
-    coordinator = entry.runtime_data.config
+    coordinator = entry.runtime_data
 
     added_entities = False
 
@@ -55,7 +55,7 @@ async def async_setup_entry(
         nonlocal added_entities
 
         if (
-            coordinator.data.configuration_control is ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control is ConfigurationControl.LOCAL
             and not added_entities
         ):
             async_add_entities(
@@ -63,7 +63,8 @@ async def async_setup_entry(
             )
             added_entities = True
         elif (
-            coordinator.data.configuration_control is not ConfigurationControl.LOCAL
+            coordinator.data.config.configuration_control
+            is not ConfigurationControl.LOCAL
             and added_entities
         ):
             entity_registry = er.async_get(hass)
@@ -82,11 +83,10 @@ class AirGradientSwitch(AirGradientEntity, SwitchEntity):
     """Defines an AirGradient switch entity."""
 
     entity_description: AirGradientSwitchEntityDescription
-    coordinator: AirGradientConfigCoordinator
 
     def __init__(
         self,
-        coordinator: AirGradientConfigCoordinator,
+        coordinator: AirGradientCoordinator,
         description: AirGradientSwitchEntityDescription,
     ) -> None:
         """Initialize AirGradient switch."""
@@ -97,7 +97,7 @@ class AirGradientSwitch(AirGradientEntity, SwitchEntity):
     @property
     def is_on(self) -> bool:
         """Return the state of the switch."""
-        return self.entity_description.value_fn(self.coordinator.data)
+        return self.entity_description.value_fn(self.coordinator.data.config)
 
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn the switch on."""

homeassistant/components/airgradient/update.py

@@ -7,7 +7,7 @@ from homeassistant.components.update import UpdateDeviceClass, UpdateEntity
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import AirGradientConfigEntry, AirGradientMeasurementCoordinator
+from . import AirGradientConfigEntry, AirGradientCoordinator
 from .entity import AirGradientEntity
 
 SCAN_INTERVAL = timedelta(hours=1)
@@ -20,18 +20,17 @@ async def async_setup_entry(
 ) -> None:
     """Set up Airgradient update platform."""
 
-    data = config_entry.runtime_data
+    coordinator = config_entry.runtime_data
 
-    async_add_entities([AirGradientUpdate(data.measurement)], True)
+    async_add_entities([AirGradientUpdate(coordinator)], True)
 
 
 class AirGradientUpdate(AirGradientEntity, UpdateEntity):
     """Representation of Airgradient Update."""
 
     _attr_device_class = UpdateDeviceClass.FIRMWARE
-    coordinator: AirGradientMeasurementCoordinator
 
-    def __init__(self, coordinator: AirGradientMeasurementCoordinator) -> None:
+    def __init__(self, coordinator: AirGradientCoordinator) -> None:
         """Initialize the entity."""
         super().__init__(coordinator)
         self._attr_unique_id = f"{coordinator.serial_number}-update"
@@ -44,7 +43,7 @@ class AirGradientUpdate(AirGradientEntity, UpdateEntity):
     @property
     def installed_version(self) -> str:
         """Return the installed version of the entity."""
-        return self.coordinator.data.firmware_version
+        return self.coordinator.data.measures.firmware_version
 
     async def async_update(self) -> None:
         """Update the entity."""

tests/components/airgradient/snapshots/test_init.ambr

@@ -57,7 +57,7 @@
     'name': 'Airgradient',
     'name_by_user': None,
     'primary_config_entry': <ANY>,
-    'serial_number': '84fce60bec38',
+    'serial_number': '84fce612f5b8',
     'suggested_area': None,
     'sw_version': '3.1.1',
     'via_device_id': None,

tests/components/airgradient/test_button.py

@@ -7,7 +7,7 @@ from airgradient import Config
 from freezegun.api import FrozenDateTimeFactory
 from syrupy import SnapshotAssertion
 
-from homeassistant.components.airgradient import DOMAIN
+from homeassistant.components.airgradient.const import DOMAIN
 from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, SERVICE_PRESS
 from homeassistant.const import ATTR_ENTITY_ID, Platform
 from homeassistant.core import HomeAssistant

tests/components/airgradient/test_config_flow.py

@@ -9,7 +9,7 @@ from airgradient import (
     ConfigurationControl,
 )
 
-from homeassistant.components.airgradient import DOMAIN
+from homeassistant.components.airgradient.const import DOMAIN
 from homeassistant.components.zeroconf import ZeroconfServiceInfo
 from homeassistant.config_entries import SOURCE_USER, SOURCE_ZEROCONF
 from homeassistant.const import CONF_HOST

tests/components/airgradient/test_init.py

@@ -4,7 +4,7 @@ from unittest.mock import AsyncMock
 
 from syrupy import SnapshotAssertion
 
-from homeassistant.components.airgradient import DOMAIN
+from homeassistant.components.airgradient.const import DOMAIN
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 

tests/components/airgradient/test_number.py

@@ -7,7 +7,7 @@ from airgradient import Config
 from freezegun.api import FrozenDateTimeFactory
 from syrupy import SnapshotAssertion
 
-from homeassistant.components.airgradient import DOMAIN
+from homeassistant.components.airgradient.const import DOMAIN
 from homeassistant.components.number import (
     ATTR_VALUE,
     DOMAIN as NUMBER_DOMAIN,

tests/components/airgradient/test_select.py

@@ -8,7 +8,7 @@ from freezegun.api import FrozenDateTimeFactory
 import pytest
 from syrupy import SnapshotAssertion
 
-from homeassistant.components.airgradient import DOMAIN
+from homeassistant.components.airgradient.const import DOMAIN
 from homeassistant.components.select import (
     DOMAIN as SELECT_DOMAIN,
     SERVICE_SELECT_OPTION,

tests/components/airgradient/test_sensor.py

@@ -8,7 +8,7 @@ from freezegun.api import FrozenDateTimeFactory
 import pytest
 from syrupy import SnapshotAssertion
 
-from homeassistant.components.airgradient import DOMAIN
+from homeassistant.components.airgradient.const import DOMAIN
 from homeassistant.const import STATE_UNAVAILABLE, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er

tests/components/airgradient/test_switch.py

@@ -7,7 +7,7 @@ from airgradient import Config
 from freezegun.api import FrozenDateTimeFactory
 from syrupy import SnapshotAssertion
 
-from homeassistant.components.airgradient import DOMAIN
+from homeassistant.components.airgradient.const import DOMAIN
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.const import (
     ATTR_ENTITY_ID,