Detect dustbin status (#34)

* report dustbin states

* add attributes

* Update vacuum.py

custom_components/purei9/purei9.py

@@ -1,6 +1,6 @@
 """Pure i9 business logic"""
 from typing import List
-from purei9_unofficial.common import BatteryStatus, RobotStates, PowerMode
+from purei9_unofficial.common import BatteryStatus, RobotStates, PowerMode, DustbinStates
 from homeassistant.components.vacuum import (
     STATE_CLEANING,
     STATE_DOCKED,
@@ -28,11 +28,18 @@
     RobotStates.Firmware_Upgrade: STATE_DOCKED
 }
 
-def state_to_hass(pure_i9_state: str, pure_i9_battery: str) -> str:
+def state_to_hass(
+    pure_i9_state: str,
+    pure_i9_battery: str,
+    purei9_dustbin: DustbinStates=DustbinStates.connected
+    ) -> str:
     """Translate Pure i9 data into a Home Assistant state constant"""
     # The Pure i9 will become "Sleeping" when docked and charged 100% OR when stopped.
     # In order to detect if it's docket or if it's just idling in the middle of a room
     # check the battery level. If it's full then we're docked.
+    if purei9_dustbin in (DustbinStates.empty, DustbinStates.full):
+        return STATE_ERROR
+
     if pure_i9_state == RobotStates.Sleeping:
         return STATE_DOCKED if pure_i9_battery == BatteryStatus.High else STATE_IDLE
 
@@ -64,6 +71,7 @@ class Params:
     available: bool = True
     firmware: str = None
     fan_speed: str = POWER_MODE_POWER
+    dustbin: DustbinStates = DustbinStates.connected
 
     def __init__(self, unique_id: str, name: str, fan_speed_list: List[str]):
         self._unique_id = unique_id

custom_components/purei9/vacuum.py

@@ -1,5 +1,5 @@
 """Home Assistant vacuum entity"""
-from typing import List, Optional, Any
+from typing import List, Optional, Any, Mapping
 from datetime import timedelta
 import homeassistant.helpers.config_validation as cv
 import voluptuous as vol
@@ -21,6 +21,7 @@
 )
 from homeassistant.const import CONF_PASSWORD, CONF_EMAIL
 from purei9_unofficial.cloudv2 import CloudClient, CloudRobot
+from purei9_unofficial.common import DustbinStates
 from . import purei9, const
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
@@ -118,6 +119,12 @@ def error(self) -> str:
         """If the vacuum reports STATE_ERROR then explain the error"""
         # According to documentation then this is required if the entity
         # can report error states. However, I can't fetch any error message.
+        if self._params.dustbin == DustbinStates.empty:
+            return "The dustbin is missing"
+
+        if self._params.dustbin == DustbinStates.full:
+            return "The dustbin needs to be emptied"
+
         return "Error"
 
     @property
@@ -136,6 +143,10 @@ def assumed_state(self) -> bool:
         return (self._assumed_next_state is not None
             or self._assumed_next_fan_speed is not None)
 
+    @property
+    def extra_state_attributes(self) -> Mapping[str, Any]:
+        return {"dustbin": self._params.dustbin.name.upper()}
+
     def start(self) -> None:
         """Start cleaning"""
         # If you click on start after clicking return, it will continue
@@ -181,12 +192,14 @@ def update(self) -> None:
         Can contain IO code.
         """
         pure_i9_battery = self._robot.getbattery()
+        purei9_dustbin = self._robot.getdustbinstatus()
 
         if self._assumed_next_state is not None:
             self._params.state = self._assumed_next_state
             self._assumed_next_state = None
         else:
-            self._params.state = purei9.state_to_hass(self._robot.getstatus(), pure_i9_battery)
+            self._params.state = purei9.state_to_hass(
+                self._robot.getstatus(), pure_i9_battery, purei9_dustbin)
 
         if self._assumed_next_fan_speed is not None:
             self._params.fan_speed = self._assumed_next_fan_speed
@@ -199,3 +212,4 @@ def update(self) -> None:
         self._params.battery = purei9.battery_to_hass(pure_i9_battery)
         self._params.available = self._robot.isconnected()
         self._params.firmware = self._robot.getfirmware()
+        self._params.dustbin = purei9_dustbin

tests/test_purei9.py

@@ -1,27 +1,30 @@
 """Test the purei9 module"""
 import unittest
-from purei9_unofficial.common import BatteryStatus, RobotStates, PowerMode
+from purei9_unofficial.common import BatteryStatus, RobotStates, PowerMode, DustbinStates
 from homeassistant.components.vacuum import (
     STATE_CLEANING,
     STATE_DOCKED,
-    STATE_IDLE
+    STATE_IDLE,
+    STATE_ERROR
 )
 from custom_components.purei9 import purei9
 
 class TestPureI9(unittest.TestCase):
     """Tests for the purei9 module"""
     data_state_to_hass = [
         # No need to test every single case. The test will become too fragile.
-        (RobotStates.Cleaning, BatteryStatus.Dead, STATE_CLEANING),
-        (RobotStates.Sleeping, BatteryStatus.High, STATE_DOCKED),
-        (RobotStates.Sleeping, BatteryStatus.Normal, STATE_IDLE)
+        (RobotStates.Cleaning, BatteryStatus.Dead, DustbinStates.connected, STATE_CLEANING),
+        (RobotStates.Sleeping, BatteryStatus.High, DustbinStates.connected, STATE_DOCKED),
+        (RobotStates.Sleeping, BatteryStatus.Normal, DustbinStates.connected, STATE_IDLE),
+        (RobotStates.Sleeping, BatteryStatus.Normal, DustbinStates.full, STATE_ERROR),
     ]
 
     def test_state_to_hass(self):
         """Test the state_to_hass function"""
-        for purei9_state, purei9_battery, expected in self.data_state_to_hass:
-            with self.subTest(purei9_state=purei9_state,purei9_battery=purei9_battery):
-                self.assertEqual(expected, purei9.state_to_hass(purei9_state, purei9_battery))
+        for purei9_state, purei9_battery, purei9_dustbin, expected in self.data_state_to_hass:
+            with self.subTest():
+                self.assertEqual(expected,
+                    purei9.state_to_hass(purei9_state, purei9_battery, purei9_dustbin))
 
     data_battery_to_hass = [
         (BatteryStatus.Dead, 0),