Move WAQI state attributes to separate sensors (#101217)

* Migrate WAQI to has entity name

* Split WAQI extra state attributes into separate sensors

* Split WAQI extra state attributes into separate sensors

* Fix test

* Support new aiowaqi

* Bump aiowaqi to 2.1.0

* Add nephelometry as possible value

* Fix test

homeassistant/components/waqi/sensor.py

@@ -1,24 +1,36 @@
 """Support for the World Air Quality Index service."""
 from __future__ import annotations
 
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
 import logging
+from typing import Any
 
-from aiowaqi import WAQIAuthenticationError, WAQIClient, WAQIConnectionError
+from aiowaqi import (
+    WAQIAirQuality,
+    WAQIAuthenticationError,
+    WAQIClient,
+    WAQIConnectionError,
+)
+from aiowaqi.models import Pollutant
 import voluptuous as vol
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
     SensorEntity,
+    SensorEntityDescription,
     SensorStateClass,
 )
 from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
 from homeassistant.const import (
-    ATTR_ATTRIBUTION,
     ATTR_TEMPERATURE,
     ATTR_TIME,
     CONF_API_KEY,
     CONF_NAME,
     CONF_TOKEN,
+    PERCENTAGE,
+    UnitOfPressure,
+    UnitOfTemperature,
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import PlatformNotReady
@@ -27,7 +39,7 @@ import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
-from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
+from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import CONF_STATION_NUMBER, DOMAIN, ISSUE_PLACEHOLDER
@@ -141,67 +153,167 @@ async def async_setup_platform(
         )
 
 
+@dataclass
+class WAQIMixin:
+    """Mixin for required keys."""
+
+    available_fn: Callable[[WAQIAirQuality], bool]
+    value_fn: Callable[[WAQIAirQuality], StateType]
+
+
+@dataclass
+class WAQISensorEntityDescription(SensorEntityDescription, WAQIMixin):
+    """Describes WAQI sensor entity."""
+
+
+SENSORS: list[WAQISensorEntityDescription] = [
+    WAQISensorEntityDescription(
+        key="air_quality",
+        device_class=SensorDeviceClass.AQI,
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.air_quality_index,
+        available_fn=lambda _: True,
+    ),
+    WAQISensorEntityDescription(
+        key="humidity",
+        device_class=SensorDeviceClass.HUMIDITY,
+        native_unit_of_measurement=PERCENTAGE,
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.humidity,
+        available_fn=lambda aq: aq.extended_air_quality.humidity is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="pressure",
+        device_class=SensorDeviceClass.PRESSURE,
+        native_unit_of_measurement=UnitOfPressure.HPA,
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.pressure,
+        available_fn=lambda aq: aq.extended_air_quality.pressure is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="temperature",
+        device_class=SensorDeviceClass.TEMPERATURE,
+        native_unit_of_measurement=UnitOfTemperature.CELSIUS,
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.temperature,
+        available_fn=lambda aq: aq.extended_air_quality.temperature is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="carbon_monoxide",
+        translation_key="carbon_monoxide",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.carbon_monoxide,
+        available_fn=lambda aq: aq.extended_air_quality.carbon_monoxide is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="nitrogen_dioxide",
+        translation_key="nitrogen_dioxide",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.nitrogen_dioxide,
+        available_fn=lambda aq: aq.extended_air_quality.nitrogen_dioxide is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="ozone",
+        translation_key="ozone",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.ozone,
+        available_fn=lambda aq: aq.extended_air_quality.ozone is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="sulphur_dioxide",
+        translation_key="sulphur_dioxide",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.sulfur_dioxide,
+        available_fn=lambda aq: aq.extended_air_quality.sulfur_dioxide is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="pm10",
+        translation_key="pm10",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.pm10,
+        available_fn=lambda aq: aq.extended_air_quality.pm10 is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="pm25",
+        translation_key="pm25",
+        state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda aq: aq.extended_air_quality.pm25,
+        available_fn=lambda aq: aq.extended_air_quality.pm25 is not None,
+    ),
+    WAQISensorEntityDescription(
+        key="dominant_pollutant",
+        translation_key="dominant_pollutant",
+        device_class=SensorDeviceClass.ENUM,
+        options=[pollutant.value for pollutant in Pollutant],
+        value_fn=lambda aq: aq.dominant_pollutant,
+        available_fn=lambda _: True,
+    ),
+]
+
+
 async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the WAQI sensor."""
     coordinator: WAQIDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities([WaqiSensor(coordinator)])
+    async_add_entities(
+        [
+            WaqiSensor(coordinator, sensor)
+            for sensor in SENSORS
+            if sensor.available_fn(coordinator.data)
+        ]
+    )
 
 
 class WaqiSensor(CoordinatorEntity[WAQIDataUpdateCoordinator], SensorEntity):
     """Implementation of a WAQI sensor."""
 
-    _attr_icon = ATTR_ICON
-    _attr_device_class = SensorDeviceClass.AQI
-    _attr_state_class = SensorStateClass.MEASUREMENT
     _attr_has_entity_name = True
-    _attr_name = None
+    entity_description: WAQISensorEntityDescription
 
-    def __init__(self, coordinator: WAQIDataUpdateCoordinator) -> None:
+    def __init__(
+        self,
+        coordinator: WAQIDataUpdateCoordinator,
+        entity_description: WAQISensorEntityDescription,
+    ) -> None:
         """Initialize the sensor."""
         super().__init__(coordinator)
-        self._attr_unique_id = f"{coordinator.data.station_id}_air_quality"
+        self.entity_description = entity_description
+        self._attr_unique_id = f"{coordinator.data.station_id}_{entity_description.key}"
         self._attr_device_info = DeviceInfo(
             identifiers={(DOMAIN, str(coordinator.data.station_id))},
             name=coordinator.data.city.name,
             entry_type=DeviceEntryType.SERVICE,
         )
+        self._attr_attribution = " and ".join(
+            attribution.name for attribution in coordinator.data.attributions
+        )
 
     @property
-    def native_value(self) -> int | None:
+    def native_value(self) -> StateType:
         """Return the state of the device."""
-        return self.coordinator.data.air_quality_index
+        return self.entity_description.value_fn(self.coordinator.data)
 
     @property
-    def extra_state_attributes(self):
+    def extra_state_attributes(self) -> Mapping[str, Any] | None:
-        """Return the state attributes of the last update."""
+        """Return old state attributes if the entity is AQI entity."""
-        attrs = {}
+        if self.entity_description.key != "air_quality":
-        try:
+            return None
-            attrs[ATTR_ATTRIBUTION] = " and ".join(
+        attrs: dict[str, Any] = {}
-                [ATTRIBUTION]
+        attrs[ATTR_TIME] = self.coordinator.data.measured_at
-                + [
+        attrs[ATTR_DOMINENTPOL] = self.coordinator.data.dominant_pollutant
-                    attribution.name
-                    for attribution in self.coordinator.data.attributions
-                ]
-            )
 
-            attrs[ATTR_TIME] = self.coordinator.data.measured_at
+        iaqi = self.coordinator.data.extended_air_quality
-            attrs[ATTR_DOMINENTPOL] = self.coordinator.data.dominant_pollutant
 
-            iaqi = self.coordinator.data.extended_air_quality
+        attribute = {
-
+            ATTR_PM2_5: iaqi.pm25,
-            attribute = {
+            ATTR_PM10: iaqi.pm10,
-                ATTR_PM2_5: iaqi.pm25,
+            ATTR_HUMIDITY: iaqi.humidity,
-                ATTR_PM10: iaqi.pm10,
+            ATTR_PRESSURE: iaqi.pressure,
-                ATTR_HUMIDITY: iaqi.humidity,
+            ATTR_TEMPERATURE: iaqi.temperature,
-                ATTR_PRESSURE: iaqi.pressure,
+            ATTR_OZONE: iaqi.ozone,
-                ATTR_TEMPERATURE: iaqi.temperature,
+            ATTR_NITROGEN_DIOXIDE: iaqi.nitrogen_dioxide,
-                ATTR_OZONE: iaqi.ozone,
+            ATTR_SULFUR_DIOXIDE: iaqi.sulfur_dioxide,
-                ATTR_NITROGEN_DIOXIDE: iaqi.nitrogen_dioxide,
+        }
-                ATTR_SULFUR_DIOXIDE: iaqi.sulfur_dioxide,
+        res_attributes = {k: v for k, v in attribute.items() if v is not None}
-            }
+        return {**attrs, **res_attributes}
-            res_attributes = {k: v for k, v in attribute.items() if v is not None}
-            return {**attrs, **res_attributes}
-        except (IndexError, KeyError):
-            return {ATTR_ATTRIBUTION: ATTRIBUTION}

homeassistant/components/waqi/strings.json

@@ -53,5 +53,39 @@
       "title": "The WAQI YAML configuration import failed",
       "description": "Configuring World Air Quality Index using YAML is being removed but there weren't any stations imported because they couldn't be found.\n\nEnsure the imported configuration is correct and remove the World Air Quality Index YAML configuration from your configuration.yaml file and continue to [set up the integration]({url}) manually."
     }
+  },
+  "entity": {
+    "sensor": {
+      "carbon_monoxide": {
+        "name": "[%key:component::sensor::entity_component::carbon_monoxide::name%]"
+      },
+      "nitrogen_dioxide": {
+        "name": "[%key:component::sensor::entity_component::nitrogen_dioxide::name%]"
+      },
+      "ozone": {
+        "name": "[%key:component::sensor::entity_component::ozone::name%]"
+      },
+      "sulphur_dioxide": {
+        "name": "[%key:component::sensor::entity_component::sulphur_dioxide::name%]"
+      },
+      "pm10": {
+        "name": "[%key:component::sensor::entity_component::pm10::name%]"
+      },
+      "pm25": {
+        "name": "[%key:component::sensor::entity_component::pm25::name%]"
+      },
+      "dominant_pollutant": {
+        "name": "Dominant pollutant",
+        "state": {
+          "co": "[%key:component::sensor::entity_component::carbon_monoxide::name%]",
+          "neph": "Nephelometry",
+          "no2": "[%key:component::sensor::entity_component::nitrogen_dioxide::name%]",
+          "o3": "[%key:component::sensor::entity_component::ozone::name%]",
+          "so2": "[%key:component::sensor::entity_component::sulphur_dioxide::name%]",
+          "pm10": "[%key:component::sensor::entity_component::pm10::name%]",
+          "pm25": "[%key:component::sensor::entity_component::pm25::name%]"
+        }
+      }
+    }
   }
 }

tests/components/waqi/fixtures/air_quality_sensor.json

@@ -23,9 +23,15 @@
     "h": {
       "v": 80
     },
+    "co": {
+      "v": 2.3
+    },
     "no2": {
       "v": 2.3
     },
+    "so2": {
+      "v": 2.3
+    },
     "o3": {
       "v": 29.4
     },

tests/components/waqi/snapshots/test_sensor.ambr

@@ -0,0 +1,181 @@
+# serializer version: 1
+# name: test_sensor
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'RIVM - Rijksinstituut voor Volksgezondheid en Milieum, Landelijk Meetnet Luchtkwaliteit and World Air Quality Index Project',
+      'device_class': 'aqi',
+      'dominentpol': <Pollutant.OZONE: 'o3'>,
+      'friendly_name': 'de Jongweg, Utrecht Air quality index',
+      'humidity': 80,
+      'nitrogen_dioxide': 2.3,
+      'ozone': 29.4,
+      'pm_10': 12,
+      'pm_2_5': 17,
+      'pressure': 1008.8,
+      'state_class': <SensorStateClass.MEASUREMENT: 'measurement'>,
+      'sulfur_dioxide': 2.3,
+      'temperature': 16,
+      'time': datetime.datetime(2023, 8, 7, 17, 0, tzinfo=datetime.timezone(datetime.timedelta(seconds=7200))),
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.de_jongweg_utrecht_air_quality_index',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '29',
+  })
+# ---
+# name: test_sensor.1
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'RIVM - Rijksinstituut voor Volksgezondheid en Milieum, Landelijk Meetnet Luchtkwaliteit and World Air Quality Index Project',
+      'device_class': 'humidity',
+      'friendly_name': 'de Jongweg, Utrecht Humidity',
+      'state_class': <SensorStateClass.MEASUREMENT: 'measurement'>,
+      'unit_of_measurement': '%',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.de_jongweg_utrecht_humidity',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '80',
+  })
+# ---
+# name: test_sensor.10
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'RIVM - Rijksinstituut voor Volksgezondheid en Milieum, Landelijk Meetnet Luchtkwaliteit and World Air Quality Index Project',
+      'device_class': 'enum',
+      'friendly_name': 'de Jongweg, Utrecht Dominant pollutant',
+      'options': list([
+        'co',
+        'no2',
+        'o3',
+        'so2',
+        'pm10',
+        'pm25',
+        'neph',
+      ]),
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.de_jongweg_utrecht_dominant_pollutant',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'o3',
+  })
+# ---
+# name: test_sensor.2
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'RIVM - Rijksinstituut voor Volksgezondheid en Milieum, Landelijk Meetnet Luchtkwaliteit and World Air Quality Index Project',
+      'device_class': 'pressure',
+      'friendly_name': 'de Jongweg, Utrecht Pressure',
+      'state_class': <SensorStateClass.MEASUREMENT: 'measurement'>,
+      'unit_of_measurement': <UnitOfPressure.HPA: 'hPa'>,
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.de_jongweg_utrecht_pressure',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '1008.8',
+  })
+# ---
+# name: test_sensor.3
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'RIVM - Rijksinstituut voor Volksgezondheid en Milieum, Landelijk Meetnet Luchtkwaliteit and World Air Quality Index Project',
+      'device_class': 'temperature',
+      'friendly_name': 'de Jongweg, Utrecht Temperature',
+      'state_class': <SensorStateClass.MEASUREMENT: 'measurement'>,
+      'unit_of_measurement': <UnitOfTemperature.CELSIUS: '°C'>,
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.de_jongweg_utrecht_temperature',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '16',
+  })
+# ---
+# name: test_sensor.4
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'RIVM - Rijksinstituut voor Volksgezondheid en Milieum, Landelijk Meetnet Luchtkwaliteit and World Air Quality Index Project',
+      'friendly_name': 'de Jongweg, Utrecht Carbon monoxide',
+      'state_class': <SensorStateClass.MEASUREMENT: 'measurement'>,
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.de_jongweg_utrecht_carbon_monoxide',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '2.3',
+  })
+# ---
+# name: test_sensor.5
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'RIVM - Rijksinstituut voor Volksgezondheid en Milieum, Landelijk Meetnet Luchtkwaliteit and World Air Quality Index Project',
+      'friendly_name': 'de Jongweg, Utrecht Nitrogen dioxide',
+      'state_class': <SensorStateClass.MEASUREMENT: 'measurement'>,
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.de_jongweg_utrecht_nitrogen_dioxide',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '2.3',
+  })
+# ---
+# name: test_sensor.6
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'RIVM - Rijksinstituut voor Volksgezondheid en Milieum, Landelijk Meetnet Luchtkwaliteit and World Air Quality Index Project',
+      'friendly_name': 'de Jongweg, Utrecht Ozone',
+      'state_class': <SensorStateClass.MEASUREMENT: 'measurement'>,
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.de_jongweg_utrecht_ozone',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '29.4',
+  })
+# ---
+# name: test_sensor.7
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'RIVM - Rijksinstituut voor Volksgezondheid en Milieum, Landelijk Meetnet Luchtkwaliteit and World Air Quality Index Project',
+      'friendly_name': 'de Jongweg, Utrecht Sulphur dioxide',
+      'state_class': <SensorStateClass.MEASUREMENT: 'measurement'>,
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.de_jongweg_utrecht_sulphur_dioxide',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '2.3',
+  })
+# ---
+# name: test_sensor.8
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'RIVM - Rijksinstituut voor Volksgezondheid en Milieum, Landelijk Meetnet Luchtkwaliteit and World Air Quality Index Project',
+      'friendly_name': 'de Jongweg, Utrecht PM10',
+      'state_class': <SensorStateClass.MEASUREMENT: 'measurement'>,
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.de_jongweg_utrecht_pm10',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '12',
+  })
+# ---
+# name: test_sensor.9
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'RIVM - Rijksinstituut voor Volksgezondheid en Milieum, Landelijk Meetnet Luchtkwaliteit and World Air Quality Index Project',
+      'friendly_name': 'de Jongweg, Utrecht PM2.5',
+      'state_class': <SensorStateClass.MEASUREMENT: 'measurement'>,
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.de_jongweg_utrecht_pm2_5',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '17',
+  })
+# ---

tests/components/waqi/test_sensor.py

@@ -3,10 +3,11 @@ import json
 from unittest.mock import patch
 
 from aiowaqi import WAQIAirQuality, WAQIError, WAQISearchResult
+from syrupy import SnapshotAssertion
 
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.components.waqi.const import CONF_STATION_NUMBER, DOMAIN
-from homeassistant.components.waqi.sensor import CONF_LOCATIONS, CONF_STATIONS
+from homeassistant.components.waqi.sensor import CONF_LOCATIONS, CONF_STATIONS, SENSORS
 from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntryState
 from homeassistant.const import (
     CONF_API_KEY,
@@ -72,7 +73,7 @@ async def test_legacy_migration_already_imported(
         assert await async_setup_component(hass, DOMAIN, {})
         await hass.async_block_till_done()
 
-    state = hass.states.get("sensor.de_jongweg_utrecht")
+    state = hass.states.get("sensor.de_jongweg_utrecht_air_quality_index")
     assert state.state == "29"
 
     hass.async_create_task(
@@ -114,13 +115,15 @@ async def test_sensor_id_migration(
     entities = er.async_entries_for_config_entry(
         entity_registry, mock_config_entry.entry_id
     )
-    assert len(entities) == 1
+    assert len(entities) == 11
     assert hass.states.get("sensor.waqi_4584")
-    assert hass.states.get("sensor.de_jongweg_utrecht") is None
+    assert hass.states.get("sensor.de_jongweg_utrecht_air_quality_index") is None
     assert entities[0].unique_id == "4584_air_quality"
 
 
-async def test_sensor(hass: HomeAssistant, mock_config_entry: MockConfigEntry) -> None:
+async def test_sensor(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry, snapshot: SnapshotAssertion
+) -> None:
     """Test failed update."""
     mock_config_entry.add_to_hass(hass)
     with patch(
@@ -131,9 +134,12 @@ async def test_sensor(hass: HomeAssistant, mock_config_entry: MockConfigEntry) -
     ):
         assert await async_setup_component(hass, DOMAIN, {})
         await hass.async_block_till_done()
-
+    entity_registry = er.async_get(hass)
-    state = hass.states.get("sensor.de_jongweg_utrecht")
+    for sensor in SENSORS:
-    assert state.state == "29"
+        entity_id = entity_registry.async_get_entity_id(
+            SENSOR_DOMAIN, DOMAIN, f"4584_{sensor.key}"
+        )
+        assert hass.states.get(entity_id) == snapshot
 
 
 async def test_updating_failed(