diff --git a/homeassistant/components/rituals_perfume_genie/sensor.py b/homeassistant/components/rituals_perfume_genie/sensor.py index c6c5c178ef5..cd637a2284f 100644 --- a/homeassistant/components/rituals_perfume_genie/sensor.py +++ b/homeassistant/components/rituals_perfume_genie/sensor.py @@ -1,7 +1,16 @@ """Support for Rituals Perfume Genie sensors.""" from __future__ import annotations -from homeassistant.components.sensor import SensorDeviceClass, SensorEntity +from collections.abc import Callable +from dataclasses import dataclass + +from pyrituals import Diffuser + +from homeassistant.components.sensor import ( + SensorDeviceClass, + SensorEntity, + SensorEntityDescription, +) from homeassistant.config_entries import ConfigEntry from homeassistant.const import PERCENTAGE, EntityCategory from homeassistant.core import HomeAssistant @@ -12,6 +21,53 @@ from .coordinator import RitualsDataUpdateCoordinator from .entity import DiffuserEntity +@dataclass +class RitualsEntityDescriptionMixin: + """Mixin values for Rituals entities.""" + + value_fn: Callable[[Diffuser], int | str] + + +@dataclass +class RitualsSensorEntityDescription( + SensorEntityDescription, RitualsEntityDescriptionMixin +): + """Class describing Rituals sensor entities.""" + + has_fn: Callable[[Diffuser], bool] = lambda _: True + + +ENTITY_DESCRIPTIONS = ( + RitualsSensorEntityDescription( + key="battery_percentage", + name="Battery", + native_unit_of_measurement=PERCENTAGE, + device_class=SensorDeviceClass.BATTERY, + value_fn=lambda diffuser: diffuser.battery_percentage, + has_fn=lambda diffuser: diffuser.has_battery, + ), + RitualsSensorEntityDescription( + key="fill", + name="Fill", + icon="mdi:beaker", + value_fn=lambda diffuser: diffuser.fill, + ), + RitualsSensorEntityDescription( + key="perfume", + name="Perfume", + icon="mdi:tag", + value_fn=lambda diffuser: diffuser.perfume, + ), + RitualsSensorEntityDescription( + key="wifi_percentage", + name="Wifi", + icon="mdi:wifi", + native_unit_of_measurement=PERCENTAGE, + value_fn=lambda diffuser: diffuser.wifi_percentage, + ), +) + + async def async_setup_entry( hass: HomeAssistant, config_entry: ConfigEntry, @@ -22,97 +78,32 @@ async def async_setup_entry( config_entry.entry_id ] - entities: list[DiffuserEntity] = [] - for coordinator in coordinators.values(): - entities.extend( - [ - DiffuserPerfumeSensor(coordinator), - DiffuserFillSensor(coordinator), - DiffuserWifiSensor(coordinator), - ] - ) - if coordinator.diffuser.has_battery: - entities.append(DiffuserBatterySensor(coordinator)) - - async_add_entities(entities) + async_add_entities( + RitualsSensorEntity(coordinator, description) + for coordinator in coordinators.values() + for description in ENTITY_DESCRIPTIONS + if description.has_fn(coordinator.diffuser) + ) -class DiffuserPerfumeSensor(DiffuserEntity, SensorEntity): - """Representation of a diffuser perfume sensor.""" +class RitualsSensorEntity(DiffuserEntity, SensorEntity): + """Representation of a diffuser sensor.""" - def __init__(self, coordinator: RitualsDataUpdateCoordinator) -> None: - """Initialize the perfume sensor.""" - super().__init__(coordinator) - self._attr_unique_id = f"{coordinator.diffuser.hublot}-perfume" - self._attr_name = f"{coordinator.diffuser.name} Perfume" - - @property - def icon(self) -> str: - """Return the perfume sensor icon.""" - if self.coordinator.diffuser.has_cartridge: - return "mdi:tag-text" - return "mdi:tag-remove" - - @property - def native_value(self) -> str: - """Return the state of the perfume sensor.""" - return self.coordinator.diffuser.perfume - - -class DiffuserFillSensor(DiffuserEntity, SensorEntity): - """Representation of a diffuser fill sensor.""" - - def __init__(self, coordinator: RitualsDataUpdateCoordinator) -> None: - """Initialize the fill sensor.""" - super().__init__(coordinator) - self._attr_unique_id = f"{coordinator.diffuser.hublot}-fill" - self._attr_name = f"{coordinator.diffuser.name} Fill" - - @property - def icon(self) -> str: - """Return the fill sensor icon.""" - if self.coordinator.diffuser.has_cartridge: - return "mdi:beaker" - return "mdi:beaker-question" - - @property - def native_value(self) -> str: - """Return the state of the fill sensor.""" - return self.coordinator.diffuser.fill - - -class DiffuserBatterySensor(DiffuserEntity, SensorEntity): - """Representation of a diffuser battery sensor.""" - - _attr_device_class = SensorDeviceClass.BATTERY - _attr_native_unit_of_measurement = PERCENTAGE + entity_description: RitualsSensorEntityDescription _attr_entity_category = EntityCategory.DIAGNOSTIC - def __init__(self, coordinator: RitualsDataUpdateCoordinator) -> None: - """Initialize the battery sensor.""" + def __init__( + self, + coordinator: RitualsDataUpdateCoordinator, + description: RitualsSensorEntityDescription, + ) -> None: + """Initialize the diffuser sensor.""" super().__init__(coordinator) - self._attr_unique_id = f"{coordinator.diffuser.hublot}-battery_percentage" - self._attr_name = f"{coordinator.diffuser.name} Battery" + self.entity_description = description + self._attr_unique_id = f"{coordinator.diffuser.hublot}-{description.key}" + self._attr_name = f"{coordinator.diffuser.name} {description.name}" @property - def native_value(self) -> int: - """Return the state of the battery sensor.""" - return self.coordinator.diffuser.battery_percentage - - -class DiffuserWifiSensor(DiffuserEntity, SensorEntity): - """Representation of a diffuser wifi sensor.""" - - _attr_native_unit_of_measurement = PERCENTAGE - _attr_entity_category = EntityCategory.DIAGNOSTIC - - def __init__(self, coordinator: RitualsDataUpdateCoordinator) -> None: - """Initialize the wifi sensor.""" - super().__init__(coordinator) - self._attr_unique_id = f"{coordinator.diffuser.hublot}-wifi_percentage" - self._attr_name = f"{coordinator.diffuser.name} Wifi" - - @property - def native_value(self) -> int: - """Return the state of the wifi sensor.""" - return self.coordinator.diffuser.wifi_percentage + def native_value(self) -> str | int: + """Return the sensor value.""" + return self.entity_description.value_fn(self.coordinator.diffuser) diff --git a/tests/components/rituals_perfume_genie/test_sensor.py b/tests/components/rituals_perfume_genie/test_sensor.py index dbbd83d1df3..af6d501d36a 100644 --- a/tests/components/rituals_perfume_genie/test_sensor.py +++ b/tests/components/rituals_perfume_genie/test_sensor.py @@ -14,7 +14,6 @@ from .common import ( init_integration, mock_config_entry, mock_diffuser_v1_battery_cartridge, - mock_diffuser_v2_no_battery_no_cartridge, ) @@ -30,7 +29,7 @@ async def test_sensors_diffuser_v1_battery_cartridge( state = hass.states.get("sensor.genie_perfume") assert state assert state.state == diffuser.perfume - assert state.attributes.get(ATTR_ICON) == "mdi:tag-text" + assert state.attributes.get(ATTR_ICON) == "mdi:tag" entry = entity_registry.async_get("sensor.genie_perfume") assert entry @@ -66,20 +65,3 @@ async def test_sensors_diffuser_v1_battery_cartridge( assert entry assert entry.unique_id == f"{hublot}-wifi_percentage" assert entry.entity_category == EntityCategory.DIAGNOSTIC - - -async def test_sensors_diffuser_v2_no_battery_no_cartridge(hass: HomeAssistant) -> None: - """Test the creation and values of the Rituals Perfume Genie sensors.""" - config_entry = mock_config_entry(unique_id="id_123_sensor_test_diffuser_v2") - - await init_integration( - hass, config_entry, [mock_diffuser_v2_no_battery_no_cartridge()] - ) - - state = hass.states.get("sensor.genie_v2_perfume") - assert state - assert state.attributes.get(ATTR_ICON) == "mdi:tag-remove" - - state = hass.states.get("sensor.genie_v2_fill") - assert state - assert state.attributes.get(ATTR_ICON) == "mdi:beaker-question"