diff --git a/homeassistant/components/tradfri/sensor.py b/homeassistant/components/tradfri/sensor.py index 3f8fa96857d..693ebeead00 100644 --- a/homeassistant/components/tradfri/sensor.py +++ b/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( - TradfriBatterySensor( - device_coordinator, - api, - gateway_id, - ) - for device_coordinator in coordinator_data[COORDINATOR_LIST] + entities: list[TradfriSensor] = [] + + for device_coordinator in coordinator_data[COORDINATOR_LIST]: + description = None 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 - _attr_native_unit_of_measurement = PERCENTAGE + entity_description: TradfriSensorEntityDescription 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) diff --git a/tests/components/tradfri/conftest.py b/tests/components/tradfri/conftest.py index f4e871d79e1..25f30237d0f 100644 --- a/tests/components/tradfri/conftest.py +++ b/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 diff --git a/tests/components/tradfri/test_fan.py b/tests/components/tradfri/test_fan.py index 4aa99f5778a..4db4ed4e585 100644 --- a/tests/components/tradfri/test_fan.py +++ b/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 = { - "fan_speed": 10, - } + state = {"fan_speed": 10, "air_quality": 12} 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 = { - "fan_speed": 10, - } + state = {"fan_speed": 10, "air_quality": 12} 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) diff --git a/tests/components/tradfri/test_sensor.py b/tests/components/tradfri/test_sensor.py index 63d4b6f84e5..04f65344125 100644 --- a/tests/components/tradfri/test_sensor.py +++ b/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}])