Add air quality sensor for Tradfri air purifier (#65070)

* Add air quality sensor for Tradfri fan platform

* Refactor, use entity description

* Fix typo

* CHange init docstring

* Let lambda handle special case

* Remove unique id

* Apply suggestions from code review

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Refactor to constants, add mixin

* Rename lambda

* Update homeassistant/components/tradfri/sensor.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Update homeassistant/components/tradfri/sensor.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Update homeassistant/components/tradfri/sensor.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Replace lambda with function

* Refactor device init

* Remove fixture scope

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/tradfri/sensor.py

@@ -2,13 +2,19 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from typing import Any
+from dataclasses import dataclass
+from typing import Any, cast
 
 from pytradfri.command import Command
+from pytradfri.device import Device
 
-from homeassistant.components.sensor import SensorDeviceClass, SensorEntity
+from homeassistant.components.sensor import (
+    SensorDeviceClass,
+    SensorEntity,
+    SensorEntityDescription,
+)
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import PERCENTAGE
+from homeassistant.const import CONCENTRATION_MICROGRAMS_PER_CUBIC_METER, PERCENTAGE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
@@ -17,6 +23,46 @@ from .const import CONF_GATEWAY_ID, COORDINATOR, COORDINATOR_LIST, DOMAIN, KEY_A
 from .coordinator import TradfriDeviceDataUpdateCoordinator
 
 
+@dataclass
+class TradfriSensorEntityDescriptionMixin:
+    """Mixin for required keys."""
+
+    value: Callable[[Device], Any | None]
+
+
+@dataclass
+class TradfriSensorEntityDescription(
+    SensorEntityDescription,
+    TradfriSensorEntityDescriptionMixin,
+):
+    """Class describing Tradfri sensor entities."""
+
+
+def _get_air_quality(device: Device) -> int | None:
+    """Fetch the air quality value."""
+    if (
+        device.air_purifier_control.air_purifiers[0].air_quality == 65535
+    ):  # The sensor returns 65535 if the fan is turned off
+        return None
+
+    return cast(int, device.air_purifier_control.air_purifiers[0].air_quality)
+
+
+SENSOR_DESCRIPTION_AQI = TradfriSensorEntityDescription(
+    device_class=SensorDeviceClass.AQI,
+    native_unit_of_measurement=CONCENTRATION_MICROGRAMS_PER_CUBIC_METER,
+    key=SensorDeviceClass.AQI,
+    value=_get_air_quality,
+)
+
+SENSOR_DESCRIPTION_BATTERY = TradfriSensorEntityDescription(
+    device_class=SensorDeviceClass.BATTERY,
+    native_unit_of_measurement=PERCENTAGE,
+    key=SensorDeviceClass.BATTERY,
+    value=lambda device: cast(int, device.device_info.battery_level),
+)
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
@@ -27,43 +73,56 @@ async def async_setup_entry(
     coordinator_data = hass.data[DOMAIN][config_entry.entry_id][COORDINATOR]
     api = coordinator_data[KEY_API]
 
-    async_add_entities(
+    entities: list[TradfriSensor] = []
-        TradfriBatterySensor(
+
-            device_coordinator,
+    for device_coordinator in coordinator_data[COORDINATOR_LIST]:
-            api,
+        description = None
-            gateway_id,
-        )
-        for device_coordinator in coordinator_data[COORDINATOR_LIST]
         if (
             not device_coordinator.device.has_light_control
             and not device_coordinator.device.has_socket_control
             and not device_coordinator.device.has_signal_repeater_control
             and not device_coordinator.device.has_air_purifier_control
-        )
+        ):
-    )
+            description = SENSOR_DESCRIPTION_BATTERY
+        elif device_coordinator.device.has_air_purifier_control:
+            description = SENSOR_DESCRIPTION_AQI
+
+        if description:
+            entities.append(
+                TradfriSensor(
+                    device_coordinator,
+                    api,
+                    gateway_id,
+                    description=description,
+                )
+            )
+
+    async_add_entities(entities)
 
 
-class TradfriBatterySensor(TradfriBaseEntity, SensorEntity):
+class TradfriSensor(TradfriBaseEntity, SensorEntity):
     """The platform class required by Home Assistant."""
 
-    _attr_device_class = SensorDeviceClass.BATTERY
+    entity_description: TradfriSensorEntityDescription
-    _attr_native_unit_of_measurement = PERCENTAGE
 
     def __init__(
         self,
         device_coordinator: TradfriDeviceDataUpdateCoordinator,
         api: Callable[[Command | list[Command]], Any],
         gateway_id: str,
+        description: TradfriSensorEntityDescription,
     ) -> None:
-        """Initialize a switch."""
+        """Initialize a Tradfri sensor."""
         super().__init__(
             device_coordinator=device_coordinator,
             api=api,
             gateway_id=gateway_id,
         )
 
+        self.entity_description = description
+
         self._refresh()  # Set initial state
 
     def _refresh(self) -> None:
         """Refresh the device."""
-        self._attr_native_value = self.coordinator.data.device_info.battery_level
+        self._attr_native_value = self.entity_description.value(self.coordinator.data)

tests/components/tradfri/conftest.py

@@ -1,5 +1,5 @@
 """Common tradfri test fixtures."""
-from unittest.mock import Mock, patch
+from unittest.mock import Mock, PropertyMock, patch
 
 import pytest
 
@@ -76,3 +76,20 @@ def mock_api_factory(mock_api):
         factory.init.return_value = factory.return_value
         factory.return_value.request = mock_api
         yield factory.return_value
+
+
+@pytest.fixture(autouse=True)
+def setup(request):
+    """
+    Set up patches for pytradfri methods for the fan platform.
+
+    This is used in test_fan as well as in test_sensor.
+    """
+    with patch(
+        "pytradfri.device.AirPurifierControl.raw",
+        new_callable=PropertyMock,
+        return_value=[{"mock": "mock"}],
+    ), patch(
+        "pytradfri.device.AirPurifierControl.air_purifiers",
+    ):
+        yield

tests/components/tradfri/test_fan.py

@@ -1,6 +1,6 @@
 """Tradfri fan (recognised as air purifiers in the IKEA ecosystem) platform tests."""
 
-from unittest.mock import MagicMock, Mock, PropertyMock, patch
+from unittest.mock import MagicMock, Mock
 
 import pytest
 from pytradfri.device import Device
@@ -10,19 +10,6 @@ from pytradfri.device.air_purifier_control import AirPurifierControl
 from .common import setup_integration
 
 
-@pytest.fixture(autouse=True, scope="module")
-def setup(request):
-    """Set up patches for pytradfri methods."""
-    with patch(
-        "pytradfri.device.AirPurifierControl.raw",
-        new_callable=PropertyMock,
-        return_value=[{"mock": "mock"}],
-    ), patch(
-        "pytradfri.device.AirPurifierControl.air_purifiers",
-    ):
-        yield
-
-
 def mock_fan(test_features=None, test_state=None, device_number=0):
     """Mock a tradfri fan/air purifier."""
     if test_features is None:
@@ -57,9 +44,7 @@ def mock_fan(test_features=None, test_state=None, device_number=0):
 
 async def test_fan(hass, mock_gateway, mock_api_factory):
     """Test that fans are correctly added."""
-    state = {
+    state = {"fan_speed": 10, "air_quality": 12}
-        "fan_speed": 10,
-    }
 
     mock_gateway.mock_devices.append(mock_fan(test_state=state))
     await setup_integration(hass)
@@ -74,9 +59,7 @@ async def test_fan(hass, mock_gateway, mock_api_factory):
 
 async def test_fan_observed(hass, mock_gateway, mock_api_factory):
     """Test that fans are correctly observed."""
-    state = {
+    state = {"fan_speed": 10, "air_quality": 12}
-        "fan_speed": 10,
-    }
 
     fan = mock_fan(test_state=state)
     mock_gateway.mock_devices.append(fan)
@@ -87,10 +70,10 @@ async def test_fan_observed(hass, mock_gateway, mock_api_factory):
 async def test_fan_available(hass, mock_gateway, mock_api_factory):
     """Test fan available property."""
 
-    fan = mock_fan(test_state={"fan_speed": 10}, device_number=1)
+    fan = mock_fan(test_state={"fan_speed": 10, "air_quality": 12}, device_number=1)
     fan.reachable = True
 
-    fan2 = mock_fan(test_state={"fan_speed": 10}, device_number=2)
+    fan2 = mock_fan(test_state={"fan_speed": 10, "air_quality": 12}, device_number=2)
     fan2.reachable = False
 
     mock_gateway.mock_devices.append(fan)
@@ -120,7 +103,7 @@ async def test_set_percentage(
 ):
     """Test setting speed of a fan."""
     # Note pytradfri style, not hass. Values not really important.
-    initial_state = {"percentage": 10, "fan_speed": 3}
+    initial_state = {"percentage": 10, "fan_speed": 3, "air_quality": 12}
     # Setup the gateway with a mock fan.
     fan = mock_fan(test_state=initial_state, device_number=0)
     mock_gateway.mock_devices.append(fan)
@@ -147,7 +130,7 @@ async def test_set_percentage(
     mock_gateway_response = responses[0]
 
     # A KeyError is raised if we don't add the 5908 response code
-    mock_gateway_response["15025"][0].update({"5908": 10})
+    mock_gateway_response["15025"][0].update({"5908": 10, "5907": 12})
 
     # Use the callback function to update the fan state.
     dev = Device(mock_gateway_response)

tests/components/tradfri/test_sensor.py

@@ -3,6 +3,7 @@
 from unittest.mock import MagicMock, Mock
 
 from .common import setup_integration
+from .test_fan import mock_fan
 
 
 def mock_sensor(test_state: list, device_number=0):
@@ -65,6 +66,20 @@ async def test_cover_battery_sensor(hass, mock_gateway, mock_api_factory):
     assert sensor_1.attributes["device_class"] == "battery"
 
 
+async def test_air_quality_sensor(hass, mock_gateway, mock_api_factory):
+    """Test that a battery sensor is correctly added."""
+    mock_gateway.mock_devices.append(
+        mock_fan(test_state={"fan_speed": 10, "air_quality": 42})
+    )
+    await setup_integration(hass)
+
+    sensor_1 = hass.states.get("sensor.tradfri_fan_0")
+    assert sensor_1 is not None
+    assert sensor_1.state == "42"
+    assert sensor_1.attributes["unit_of_measurement"] == "µg/m³"
+    assert sensor_1.attributes["device_class"] == "aqi"
+
+
 async def test_sensor_observed(hass, mock_gateway, mock_api_factory):
     """Test that sensors are correctly observed."""
     sensor = mock_sensor(test_state=[{"attribute": "battery_level", "value": 60}])