From b452db8b67188a36958a38a5267804afe0947aed Mon Sep 17 00:00:00 2001
From: Xiaonan Shen <s@sxn.dev>
Date: Tue, 19 May 2020 22:19:27 -0700
Subject: [PATCH] Add dynamic icon to roomba battery sensor (#35647)

* Add dynamic icon to roomba battery sensor

* Fix lint

* Rename _state to _robot_state
---
 .../components/roomba/irobot_base.py          | 59 ++++++++++---------
 homeassistant/components/roomba/sensor.py     | 18 ++++--
 2 files changed, 44 insertions(+), 33 deletions(-)

diff --git a/homeassistant/components/roomba/irobot_base.py b/homeassistant/components/roomba/irobot_base.py
index f510f4965b0..035428192c0 100644
--- a/homeassistant/components/roomba/irobot_base.py
+++ b/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._name = vacuum_state.get("name")
-        self._version = vacuum_state.get("softwareVer")
-        self._sku = vacuum_state.get("sku")
+        self.vacuum_state = roomba_reported_state(roomba)
+        self._name = self.vacuum_state.get("name")
+        self._version = self.vacuum_state.get("softwareVer")
+        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):
-        """Filter the new state."""
-        raise NotImplementedError
+    def new_state_filter(self, new_state):  # pylint: disable=no-self-use
+        """Filter out wifi state messages."""
+        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", {})
-        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
+        return self._robot_state
 
     @property
     def available(self) -> bool:
@@ -215,14 +224,10 @@ class IRobotVacuum(IRobotEntity, StateVacuumEntity):
 
     def on_message(self, json_data):
         """Update state on message change."""
-        new_state = json_data.get("state", {}).get("reported", {})
-        if (
-            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()
+        state = json_data.get("state", {}).get("reported", {})
+        if self.new_state_filter(state):
+            _LOGGER.debug("Got new state from the vacuum: %s", json_data)
+            self.schedule_update_ha_state()
 
     async def async_start(self):
         """Start or resume the cleaning task."""
diff --git a/homeassistant/components/roomba/sensor.py b/homeassistant/components/roomba/sensor.py
index 40dbd52e158..435f26510f8 100644
--- a/homeassistant/components/roomba/sensor.py
+++ b/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")
-
-    def new_state_filter(self, new_state):
-        """Filter the new state."""
-        return "batPct" in new_state
+        return self._battery_level