Add dynamic icon to roomba battery sensor (#35647)

* Add dynamic icon to roomba battery sensor

* Fix lint

* Rename _state to _robot_state

homeassistant/components/roomba/irobot_base.py

@@ -68,10 +68,10 @@ class IRobotEntity(Entity):
         """Initialize the iRobot handler."""
         self.vacuum = roomba
         self._blid = blid
-        vacuum_state = roomba_reported_state(roomba)
+        self.vacuum_state = roomba_reported_state(roomba)
-        self._name = vacuum_state.get("name")
+        self._name = self.vacuum_state.get("name")
-        self._version = vacuum_state.get("softwareVer")
+        self._version = self.vacuum_state.get("softwareVer")
-        self._sku = vacuum_state.get("sku")
+        self._sku = self.vacuum_state.get("sku")
 
     @property
     def should_poll(self):
@@ -99,13 +99,32 @@ class IRobotEntity(Entity):
             "model": self._sku,
         }
 
+    @property
+    def _battery_level(self):
+        """Return the battery level of the vacuum cleaner."""
+        return self.vacuum_state.get("batPct")
+
+    @property
+    def _robot_state(self):
+        """Return the state of the vacuum cleaner."""
+        clean_mission_status = self.vacuum_state.get("cleanMissionStatus", {})
+        cycle = clean_mission_status.get("cycle")
+        phase = clean_mission_status.get("phase")
+        try:
+            state = STATE_MAP[phase]
+        except KeyError:
+            return STATE_ERROR
+        if cycle != "none" and state in (STATE_IDLE, STATE_DOCKED):
+            state = STATE_PAUSED
+        return state
+
     async def async_added_to_hass(self):
         """Register callback function."""
         self.vacuum.register_on_message_callback(self.on_message)
 
-    def new_state_filter(self, new_state):
+    def new_state_filter(self, new_state):  # pylint: disable=no-self-use
-        """Filter the new state."""
+        """Filter out wifi state messages."""
-        raise NotImplementedError
+        return len(new_state) > 1 or "signal" not in new_state
 
     def on_message(self, json_data):
         """Update state on message change."""
@@ -120,7 +139,6 @@ class IRobotVacuum(IRobotEntity, StateVacuumEntity):
     def __init__(self, roomba, blid):
         """Initialize the iRobot handler."""
         super().__init__(roomba, blid)
-        self.vacuum_state = roomba_reported_state(roomba)
         self._cap_position = self.vacuum_state.get("cap", {}).get("pose") == 1
 
     @property
@@ -141,21 +159,12 @@ class IRobotVacuum(IRobotEntity, StateVacuumEntity):
     @property
     def battery_level(self):
         """Return the battery level of the vacuum cleaner."""
-        return self.vacuum_state.get("batPct")
+        return self._battery_level
 
     @property
     def state(self):
         """Return the state of the vacuum cleaner."""
-        clean_mission_status = self.vacuum_state.get("cleanMissionStatus", {})
+        return self._robot_state
-        cycle = clean_mission_status.get("cycle")
-        phase = clean_mission_status.get("phase")
-        try:
-            state = STATE_MAP[phase]
-        except KeyError:
-            return STATE_ERROR
-        if cycle != "none" and state in (STATE_IDLE, STATE_DOCKED):
-            state = STATE_PAUSED
-        return state
 
     @property
     def available(self) -> bool:
@@ -215,13 +224,9 @@ class IRobotVacuum(IRobotEntity, StateVacuumEntity):
 
     def on_message(self, json_data):
         """Update state on message change."""
-        new_state = json_data.get("state", {}).get("reported", {})
+        state = json_data.get("state", {}).get("reported", {})
-        if (
+        if self.new_state_filter(state):
-            len(new_state) == 1 and "signal" in new_state
-        ):  # filter out wifi stat messages
-            return
             _LOGGER.debug("Got new state from the vacuum: %s", json_data)
-        self.vacuum_state = roomba_reported_state(self.vacuum)
             self.schedule_update_ha_state()
 
     async def async_start(self):

homeassistant/components/roomba/sensor.py

@@ -1,9 +1,10 @@
 """Sensor for checking the battery level of Roomba."""
 import logging
 
+from homeassistant.components.vacuum import STATE_DOCKED
 from homeassistant.const import DEVICE_CLASS_BATTERY, UNIT_PERCENTAGE
+from homeassistant.helpers.icon import icon_for_battery_level
 
-from . import roomba_reported_state
 from .const import BLID, DOMAIN, ROOMBA_SESSION
 from .irobot_base import IRobotEntity
 
@@ -42,11 +43,16 @@ class RoombaBattery(IRobotEntity):
         """Return the unit_of_measurement of the device."""
         return UNIT_PERCENTAGE
 
+    @property
+    def icon(self):
+        """Return the icon for the battery."""
+        charging = bool(self._robot_state == STATE_DOCKED)
+
+        return icon_for_battery_level(
+            battery_level=self._battery_level, charging=charging
+        )
+
     @property
     def state(self):
         """Return the state of the sensor."""
-        return roomba_reported_state(self.vacuum).get("batPct")
+        return self._battery_level
-
-    def new_state_filter(self, new_state):
-        """Filter the new state."""
-        return "batPct" in new_state