From 26cbc26af180a9fb948b8e6f221f6a20f9a0f650 Mon Sep 17 00:00:00 2001
From: Luke Bonaccorsi <luke@lukeb.co.uk>
Date: Thu, 17 Aug 2023 17:00:40 +0100
Subject: [PATCH] Separate consumables into it's own feature

---
 custom_components/robovac/robovac.py | 43 ++++++++++++++++++++++------
 custom_components/robovac/vacuum.py  | 21 ++++++--------
 2 files changed, 43 insertions(+), 21 deletions(-)

diff --git a/custom_components/robovac/robovac.py b/custom_components/robovac/robovac.py
index 9b12139..6500b30 100644
--- a/custom_components/robovac/robovac.py
+++ b/custom_components/robovac/robovac.py
@@ -18,6 +18,7 @@ class RoboVacEntityFeature(IntEnum):
     MAP = 512
     BOOST_IQ = 1024
 
+
 ROBOVAC_SERIES = {
     "C": [
         "T2103",
@@ -28,6 +29,7 @@ ROBOVAC_SERIES = {
         "T2123",
         "T2128",
         "T2130",
+        "T2132",
     ],
     "G": [
         "T1250",
@@ -37,37 +39,54 @@ ROBOVAC_SERIES = {
         "T2253",
         "T2150",
         "T2255",
+        "T2256",
+        "T2257",
+        "T2258",
+        "T2259",
+        "T2270",
+        "T2272",
+        "T2273",
     ],
-    "L": ["T2182", "T2192"],
+    "L": ["T2181", "T2182", "T2190", "T2192", "T2193", "T2194"],
     "X": ["T2261", "T2262", "T2320"],
 }
 
 HAS_MAP_FEATURE = ["T2253", *ROBOVAC_SERIES["L"], *ROBOVAC_SERIES["X"]]
 
+HAS_CONSUMABLES = [
+    "T1250",
+    "T2181",
+    "T2182",
+    "T2190",
+    "T2193",
+    "T2194",
+    "T2253",
+    "T2256",
+    "T2258",
+    "T2261",
+    "T2273",
+    "T2320",
+]
+
 ROBOVAC_SERIES_FEATURES = {
     "C": RoboVacEntityFeature.EDGE | RoboVacEntityFeature.SMALL_ROOM,
     "G": RoboVacEntityFeature.CLEANING_TIME
     | RoboVacEntityFeature.CLEANING_AREA
     | RoboVacEntityFeature.DO_NOT_DISTURB
-    | RoboVacEntityFeature.AUTO_RETURN
-    | RoboVacEntityFeature.CONSUMABLES,
+    | RoboVacEntityFeature.AUTO_RETURN,
     "L": RoboVacEntityFeature.CLEANING_TIME
     | RoboVacEntityFeature.CLEANING_AREA
     | RoboVacEntityFeature.DO_NOT_DISTURB
     | RoboVacEntityFeature.AUTO_RETURN
-    | RoboVacEntityFeature.CONSUMABLES
     | RoboVacEntityFeature.ROOM
     | RoboVacEntityFeature.ZONE
-    | RoboVacEntityFeature.MAP
     | RoboVacEntityFeature.BOOST_IQ,
     "X": RoboVacEntityFeature.CLEANING_TIME
     | RoboVacEntityFeature.CLEANING_AREA
     | RoboVacEntityFeature.DO_NOT_DISTURB
     | RoboVacEntityFeature.AUTO_RETURN
-    | RoboVacEntityFeature.CONSUMABLES
     | RoboVacEntityFeature.ROOM
     | RoboVacEntityFeature.ZONE
-    | RoboVacEntityFeature.MAP
     | RoboVacEntityFeature.BOOST_IQ,
 }
 
@@ -120,7 +139,15 @@ class RoboVac(TuyaDevice):
         return supportedFeatures
 
     def getRoboVacFeatures(self):
-        return ROBOVAC_SERIES_FEATURES[self.getRoboVacSeries()]
+        supportedFeatures = ROBOVAC_SERIES_FEATURES[self.getRoboVacSeries()]
+
+        if self.model_code in HAS_MAP_FEATURE:
+            supportedFeatures |= RoboVacEntityFeature.MAP
+
+        if self.model_code in HAS_CONSUMABLES:
+            supportedFeatures |= RoboVacEntityFeature.CONSUMABLES
+
+        return supportedFeatures
 
     def getRoboVacSeries(self):
         for series, models in ROBOVAC_SERIES.items():
diff --git a/custom_components/robovac/vacuum.py b/custom_components/robovac/vacuum.py
index 4b65ece..05a07f4 100644
--- a/custom_components/robovac/vacuum.py
+++ b/custom_components/robovac/vacuum.py
@@ -284,22 +284,17 @@ class RoboVacEntity(StateVacuumEntity):
         self._attr_cleaning_time = self.tuyastatus.get(TUYA_CODES.CLEANING_TIME)
         self._attr_auto_return = self.tuyastatus.get(TUYA_CODES.AUTO_RETURN)
         self._attr_do_not_disturb = self.tuyastatus.get(TUYA_CODES.DO_NOT_DISTURB)
-        if self.tuyastatus.get(TUYA_CODES.G_CONSUMABLES) is not None:
-            self._attr_consumables = ast.literal_eval(
-                base64.b64decode(self.tuyastatus.get(TUYA_CODES.G_CONSUMABLES)).decode(
-                    "ascii"
-                )
-            )["consumable"]["duration"]
-        # For X8
         self._attr_boost_iq = self.tuyastatus.get(TUYA_CODES.BOOST_IQ)
         # self.map_data = self.tuyastatus.get("121")
         # self.erro_msg? = self.tuyastatus.get("124")
-        if self.tuyastatus.get(TUYA_CODES.X_CONSUMABLES) is not None:
-            self._attr_consumables = ast.literal_eval(
-                base64.b64decode(self.tuyastatus.get(TUYA_CODES.X_CONSUMABLES)).decode(
-                    "ascii"
-                )
-            )["consumable"]["duration"]
+        if self.robovac_supported & RoboVacEntityFeature.CONSUMABLES:
+            robovac_series = self.vacuum.getRoboVacSeries()
+            if self.tuyastatus.get(TUYA_CODES["{}_CONSUMABLES".format(robovac_series)]) is not None:
+                self._attr_consumables = ast.literal_eval(
+                    base64.b64decode(self.tuyastatus.get(TUYA_CODES["{}_CONSUMABLES".format(robovac_series)])).decode(
+                        "ascii"
+                    )
+                )["consumable"]["duration"]
 
     async def async_locate(self, **kwargs):
         """Locate the vacuum cleaner."""