Use entity descriptions in Roomba (#102323)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/roomba/irobot_base.py

@@ -101,20 +101,23 @@ class IRobotEntity(Entity):
         )
 
     @property
-    def _battery_level(self):
+    def battery_level(self):
         """Return the battery level of the vacuum cleaner."""
         return self.vacuum_state.get("batPct")
 
     @property
-    def _run_stats(self):
+    def run_stats(self):
+        """Return the run stats."""
         return self.vacuum_state.get("bbrun")
 
     @property
-    def _mission_stats(self):
+    def mission_stats(self):
+        """Return the mission stats."""
         return self.vacuum_state.get("bbmssn")
 
     @property
-    def _battery_stats(self):
+    def battery_stats(self):
+        """Return the battery stats."""
         return self.vacuum_state.get("bbchg3")
 
     @property
@@ -158,11 +161,6 @@ class IRobotVacuum(IRobotEntity, StateVacuumEntity):
         super().__init__(roomba, blid)
         self._cap_position = self.vacuum_state.get("cap", {}).get("pose") == 1
 
-    @property
-    def battery_level(self):
-        """Return the battery level of the vacuum cleaner."""
-        return self._battery_level
-
     @property
     def state(self):
         """Return the state of the vacuum cleaner."""

homeassistant/components/roomba/sensor.py

@@ -1,18 +1,121 @@
-"""Sensor for checking the battery level of Roomba."""
+"""Sensor platform for Roomba."""
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from roombapy import Roomba
+
 from homeassistant.components.sensor import (
     SensorDeviceClass,
     SensorEntity,
+    SensorEntityDescription,
     SensorStateClass,
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE, EntityCategory, UnitOfTime
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.typing import StateType
 
 from .const import BLID, DOMAIN, ROOMBA_SESSION
 from .irobot_base import IRobotEntity
 
 
+@dataclass
+class RoombaSensorEntityDescriptionMixin:
+    """Mixin for describing Roomba data."""
+
+    value_fn: Callable[[IRobotEntity], StateType]
+
+
+@dataclass
+class RoombaSensorEntityDescription(
+    SensorEntityDescription, RoombaSensorEntityDescriptionMixin
+):
+    """Immutable class for describing Roomba data."""
+
+
+SENSORS: list[RoombaSensorEntityDescription] = [
+    RoombaSensorEntityDescription(
+        key="battery",
+        native_unit_of_measurement=PERCENTAGE,
+        device_class=SensorDeviceClass.BATTERY,
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda self: self.battery_level,
+    ),
+    RoombaSensorEntityDescription(
+        key="battery_cycles",
+        translation_key="battery_cycles",
+        state_class=SensorStateClass.MEASUREMENT,
+        icon="mdi:counter",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda self: self.battery_stats.get("nLithChrg")
+        or self.battery_stats.get("nNimhChrg"),
+    ),
+    RoombaSensorEntityDescription(
+        key="total_cleaning_time",
+        translation_key="total_cleaning_time",
+        icon="mdi:clock",
+        native_unit_of_measurement=UnitOfTime.HOURS,
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda self: self.run_stats.get("hr"),
+    ),
+    RoombaSensorEntityDescription(
+        key="average_mission_time",
+        translation_key="average_mission_time",
+        icon="mdi:clock",
+        native_unit_of_measurement=UnitOfTime.MINUTES,
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda self: self.mission_stats.get("aMssnM"),
+    ),
+    RoombaSensorEntityDescription(
+        key="total_missions",
+        translation_key="total_missions",
+        icon="mdi:counter",
+        state_class=SensorStateClass.MEASUREMENT,
+        native_unit_of_measurement="Missions",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda self: self.mission_stats.get("nMssn"),
+    ),
+    RoombaSensorEntityDescription(
+        key="successful_missions",
+        translation_key="successful_missions",
+        icon="mdi:counter",
+        state_class=SensorStateClass.MEASUREMENT,
+        native_unit_of_measurement="Missions",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda self: self.mission_stats.get("nMssnOk"),
+    ),
+    RoombaSensorEntityDescription(
+        key="canceled_missions",
+        translation_key="canceled_missions",
+        icon="mdi:counter",
+        state_class=SensorStateClass.MEASUREMENT,
+        native_unit_of_measurement="Missions",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda self: self.mission_stats.get("nMssnC"),
+    ),
+    RoombaSensorEntityDescription(
+        key="failed_missions",
+        translation_key="failed_missions",
+        icon="mdi:counter",
+        state_class=SensorStateClass.MEASUREMENT,
+        native_unit_of_measurement="Missions",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda self: self.mission_stats.get("nMssnF"),
+    ),
+    RoombaSensorEntityDescription(
+        key="scrubs_count",
+        translation_key="scrubs",
+        icon="mdi:counter",
+        state_class=SensorStateClass.MEASUREMENT,
+        native_unit_of_measurement="Scrubs",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        value_fn=lambda self: self.run_stats.get("nScrubs"),
+        entity_registry_enabled_default=False,
+    ),
+]
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
@@ -23,188 +126,32 @@ async def async_setup_entry(
     roomba = domain_data[ROOMBA_SESSION]
     blid = domain_data[BLID]
 
-    roomba_vac = RoombaBattery(roomba, blid)
-    roomba_battery_cycles = BatteryCycles(roomba, blid)
-    roomba_cleaning_time = CleaningTime(roomba, blid)
-    roomba_average_mission_time = AverageMissionTime(roomba, blid)
-    roomba_total_missions = MissionSensor(roomba, blid, "total", "nMssn")
-    roomba_success_missions = MissionSensor(roomba, blid, "successful", "nMssnOk")
-    roomba_canceled_missions = MissionSensor(roomba, blid, "canceled", "nMssnC")
-    roomba_failed_missions = MissionSensor(roomba, blid, "failed", "nMssnF")
-    roomba_scrubs_count = ScrubsCount(roomba, blid)
-
     async_add_entities(
-        [
+        RoombaSensor(roomba, blid, entity_description) for entity_description in SENSORS
-            roomba_vac,
-            roomba_battery_cycles,
-            roomba_cleaning_time,
-            roomba_average_mission_time,
-            roomba_total_missions,
-            roomba_success_missions,
-            roomba_canceled_missions,
-            roomba_failed_missions,
-            roomba_scrubs_count,
-        ],
     )
 
 
-class RoombaBattery(IRobotEntity, SensorEntity):
+class RoombaSensor(IRobotEntity, SensorEntity):
-    """Class to hold Roomba Sensor basic info."""
+    """Roomba sensor."""
 
-    _attr_entity_category = EntityCategory.DIAGNOSTIC
+    entity_description: RoombaSensorEntityDescription
-    _attr_device_class = SensorDeviceClass.BATTERY
-    _attr_native_unit_of_measurement = PERCENTAGE
 
-    @property
+    def __init__(
-    def unique_id(self):
+        self,
-        """Return the ID of this sensor."""
+        roomba: Roomba,
-        return f"battery_{self._blid}"
+        blid: str,
-
+        entity_description: RoombaSensorEntityDescription,
-    @property
+    ) -> None:
-    def native_value(self):
+        """Initialize Roomba sensor."""
-        """Return the state of the sensor."""
-        return self._battery_level
-
-
-class BatteryCycles(IRobotEntity, SensorEntity):
-    """Class to hold Roomba Sensor basic info."""
-
-    _attr_entity_category = EntityCategory.DIAGNOSTIC
-    _attr_icon = "mdi:counter"
-
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return "Battery cycles"
-
-    @property
-    def unique_id(self):
-        """Return the ID of this sensor."""
-        return f"battery_cycles_{self._blid}"
-
-    @property
-    def state_class(self):
-        """Return the state class of this entity, from STATE_CLASSES, if any."""
-        return SensorStateClass.MEASUREMENT
-
-    @property
-    def native_value(self):
-        """Return the state of the sensor."""
-        return self._battery_stats.get("nLithChrg") or self._battery_stats.get(
-            "nNimhChrg"
-        )
-
-
-class CleaningTime(IRobotEntity, SensorEntity):
-    """Class to hold Roomba Sensor basic info."""
-
-    _attr_entity_category = EntityCategory.DIAGNOSTIC
-    _attr_icon = "mdi:clock"
-    _attr_native_unit_of_measurement = UnitOfTime.HOURS
-
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return f"{self._name} cleaning time total"
-
-    @property
-    def unique_id(self):
-        """Return the ID of this sensor."""
-        return f"total_cleaning_time_{self._blid}"
-
-    @property
-    def native_value(self):
-        """Return the state of the sensor."""
-        return self._run_stats.get("hr")
-
-
-class AverageMissionTime(IRobotEntity, SensorEntity):
-    """Class to hold Roomba Sensor basic info."""
-
-    _attr_entity_category = EntityCategory.DIAGNOSTIC
-    _attr_icon = "mdi:clock"
-
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return "Average mission time"
-
-    @property
-    def unique_id(self):
-        """Return the ID of this sensor."""
-        return f"average_mission_time_{self._blid}"
-
-    @property
-    def native_unit_of_measurement(self):
-        """Return the unit_of_measurement of the device."""
-        return UnitOfTime.MINUTES
-
-    @property
-    def native_value(self):
-        """Return the state of the sensor."""
-        return self._mission_stats.get("aMssnM")
-
-
-class MissionSensor(IRobotEntity, SensorEntity):
-    """Class to hold the Roomba missions info."""
-
-    def __init__(self, roomba, blid, mission_type, mission_value_string):
-        """Initialise iRobot sensor with mission details."""
         super().__init__(roomba, blid)
-        self._mission_type = mission_type
+        self.entity_description = entity_description
-        self._mission_value_string = mission_value_string
-
-    _attr_entity_category = EntityCategory.DIAGNOSTIC
-    _attr_icon = "mdi:counter"
 
     @property
-    def name(self):
+    def unique_id(self) -> str:
-        """Return the name of the sensor."""
+        """Return a unique ID."""
-        return f"Missions {self._mission_type}"
+        return f"{self.entity_description.key}_{self._blid}"
 
     @property
-    def unique_id(self):
+    def native_value(self) -> StateType:
-        """Return the ID of this sensor."""
-        return f"{self._mission_type}_missions_{self._blid}"
-
-    @property
-    def state_class(self):
-        """Return the state class of this entity, from STATE_CLASSES, if any."""
-        return SensorStateClass.MEASUREMENT
-
-    @property
-    def native_value(self):
         """Return the state of the sensor."""
-        return self._mission_stats.get(self._mission_value_string)
+        return self.entity_description.value_fn(self)
-
-
-class ScrubsCount(IRobotEntity, SensorEntity):
-    """Class to hold Roomba Sensor basic info."""
-
-    _attr_entity_category = EntityCategory.DIAGNOSTIC
-    _attr_icon = "mdi:counter"
-
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return "Scrubs count"
-
-    @property
-    def unique_id(self):
-        """Return the ID of this sensor."""
-        return f"scrubs_count_{self._blid}"
-
-    @property
-    def state_class(self):
-        """Return the state class of this entity, from STATE_CLASSES, if any."""
-        return SensorStateClass.MEASUREMENT
-
-    @property
-    def entity_registry_enabled_default(self):
-        """Disable sensor by default."""
-        return False
-
-    @property
-    def native_value(self):
-        """Return the state of the sensor."""
-        return self._run_stats.get("nScrubs")

homeassistant/components/roomba/strings.json

@@ -53,6 +53,32 @@
       "bin_full": {
         "name": "Bin full"
       }
+    },
+    "sensor": {
+      "battery_cycles": {
+        "name": "Battery cycles"
+      },
+      "total_cleaning_time": {
+        "name": "Total cleaning time"
+      },
+      "average_mission_time": {
+        "name": "Average mission time"
+      },
+      "total_missions": {
+        "name": "Total missions"
+      },
+      "successful_missions": {
+        "name": "Successful missions"
+      },
+      "canceled_missions": {
+        "name": "Canceled missions"
+      },
+      "failed_missions": {
+        "name": "Failed missions"
+      },
+      "scrubs_count": {
+        "name": "Scrubs"
+      }
     }
   }
 }