Add entity descriptions to sensors of Rituals Perfume Genie

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,57 @@
 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,
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda diffuser: diffuser.battery_percentage,
+        has_fn=lambda diffuser: diffuser.has_battery,
+    ),
+    RitualsSensorEntityDescription(
+        key="fill",
+        name="Fill",
+        icon="mdi:beaker",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda diffuser: diffuser.fill,
+    ),
+    RitualsSensorEntityDescription(
+        key="perfume",
+        name="Perfume",
+        icon="mdi:tag",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda diffuser: diffuser.perfume,
+    ),
+    RitualsSensorEntityDescription(
+        key="wifi_percentage",
+        name="Wifi",
+        icon="mdi:wifi",
+        native_unit_of_measurement=PERCENTAGE,
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda diffuser: diffuser.wifi_percentage,
+    ),
+)
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
@@ -22,97 +82,31 @@ 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)
-
-
-class DiffuserPerfumeSensor(DiffuserEntity, SensorEntity):
-    """Representation of a diffuser perfume 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
-    _attr_entity_category = EntityCategory.DIAGNOSTIC
-
-    def __init__(self, coordinator: RitualsDataUpdateCoordinator) -> None:
-        """Initialize the battery sensor."""
-        super().__init__(coordinator)
-        self._attr_unique_id = f"{coordinator.diffuser.hublot}-battery_percentage"
-        self._attr_name = f"{coordinator.diffuser.name} Battery"
-
-    @property
-    def native_value(self) -> int:
-        """Return the state of the battery sensor."""
-        return self.coordinator.diffuser.battery_percentage
+    async_add_entities(
+        RitualsSensorEntity(coordinator, description)
+        for coordinator in coordinators.values()
+        for description in ENTITY_DESCRIPTIONS
+        if description.has_fn(coordinator.diffuser)
+    )
 
 
-class DiffuserWifiSensor(DiffuserEntity, SensorEntity):
-    """Representation of a diffuser wifi sensor."""
+class RitualsSensorEntity(DiffuserEntity, SensorEntity):
+    """Representation of a diffuser sensor."""
 
-    _attr_native_unit_of_measurement = PERCENTAGE
-    _attr_entity_category = EntityCategory.DIAGNOSTIC
+    entity_description: RitualsSensorEntityDescription
 
-    def __init__(self, coordinator: RitualsDataUpdateCoordinator) -> None:
-        """Initialize the wifi sensor."""
+    def __init__(
+        self,
+        coordinator: RitualsDataUpdateCoordinator,
+        description: RitualsSensorEntityDescription,
+    ) -> None:
+        """Initialize the diffuser sensor."""
         super().__init__(coordinator)
-        self._attr_unique_id = f"{coordinator.diffuser.hublot}-wifi_percentage"
-        self._attr_name = f"{coordinator.diffuser.name} Wifi"
+        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 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)

tests/components/rituals_perfume_genie/test_sensor.py

@@ -14,7 +14,6 @@
     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"