From 3625657886bfafb8468d0bbbf7fa1a4a76fb2bf5 Mon Sep 17 00:00:00 2001
From: Luke Bonaccorsi <luke@lukeb.co.uk>
Date: Mon, 4 Sep 2023 12:37:58 +0100
Subject: [PATCH] Try various consumables codes rather than just relying on
 fixed series ones

---
 custom_components/robovac/vacuum.py | 23 ++++++++++-------------
 1 file changed, 10 insertions(+), 13 deletions(-)

diff --git a/custom_components/robovac/vacuum.py b/custom_components/robovac/vacuum.py
index 14a8535..5d1a727 100644
--- a/custom_components/robovac/vacuum.py
+++ b/custom_components/robovac/vacuum.py
@@ -98,8 +98,7 @@ class TUYA_CODES(StrEnum):
     AUTO_RETURN = "135"
     DO_NOT_DISTURB = "107"
     BOOST_IQ = "118"
-    G_CONSUMABLES = "142"
-    X_CONSUMABLES = "116"
+    CONSUMABLES = ["142", "116"]
 
 
 async def async_setup_entry(
@@ -318,18 +317,16 @@ class RoboVacEntity(StateVacuumEntity):
         # self.map_data = self.tuyastatus.get("121")
         # self.erro_msg? = self.tuyastatus.get("124")
         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)]
+            for CONSUMABLE_CODE in TUYA_CODES.CONSUMABLES:
+                if (
+                    CONSUMABLE_CODE in self.tuyastatus
+                    and self.tuyastatus.get(CONSUMABLE_CODE) is not None
+                ):
+                    self._attr_consumables = ast.literal_eval(
+                        base64.b64decode(self.tuyastatus.get(CONSUMABLE_CODE)).decode(
+                            "ascii"
                         )
-                    ).decode("ascii")
-                )["consumable"]["duration"]
+                    )["consumable"]["duration"]
 
     async def async_locate(self, **kwargs):
         """Locate the vacuum cleaner."""