wip

Author: CodeFoodPixels

custom_components/robovac/robovac.py

@@ -1,107 +1,6 @@
-from enum import IntEnum
-from homeassistant.components.vacuum import VacuumEntityFeature
+from .vacuums.base import RobovacCommand
 from .tuyalocalapi import TuyaDevice
-
-
-class RoboVacEntityFeature(IntEnum):
-    """Supported features of the RoboVac entity."""
-
-    EDGE = 1
-    SMALL_ROOM = 2
-    CLEANING_TIME = 4
-    CLEANING_AREA = 8
-    DO_NOT_DISTURB = 16
-    AUTO_RETURN = 32
-    CONSUMABLES = 64
-    ROOM = 128
-    ZONE = 256
-    MAP = 512
-    BOOST_IQ = 1024
-
-
-ROBOVAC_SERIES = {
-    "C": [
-        "T2103",
-        "T2117",
-        "T2118",
-        "T2119",
-        "T2120",
-        "T2123",
-        "T2128",
-        "T2130",
-        "T2132",
-    ],
-    "G": [
-        "T1250",
-        "T2250",
-        "T2251",
-        "T2252",
-        "T2253",
-        "T2254",
-        "T2150",
-        "T2255",
-        "T2256",
-        "T2257",
-        "T2258",
-        "T2259",
-        "T2270",
-        "T2272",
-        "T2273",
-    ],
-    "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,
-    "L": RoboVacEntityFeature.CLEANING_TIME
-    | RoboVacEntityFeature.CLEANING_AREA
-    | RoboVacEntityFeature.DO_NOT_DISTURB
-    | RoboVacEntityFeature.AUTO_RETURN
-    | RoboVacEntityFeature.ROOM
-    | RoboVacEntityFeature.ZONE
-    | RoboVacEntityFeature.BOOST_IQ,
-    "X": RoboVacEntityFeature.CLEANING_TIME
-    | RoboVacEntityFeature.CLEANING_AREA
-    | RoboVacEntityFeature.DO_NOT_DISTURB
-    | RoboVacEntityFeature.AUTO_RETURN
-    | RoboVacEntityFeature.ROOM
-    | RoboVacEntityFeature.ZONE
-    | RoboVacEntityFeature.BOOST_IQ,
-}
-
-ROBOVAC_SERIES_FAN_SPEEDS = {
-    "C": ["No Suction", "Standard", "Boost IQ", "Max"],
-    "G": ["Standard", "Turbo", "Max", "Boost IQ"],
-    "L": ["Quiet", "Standard", "Turbo", "Max"],
-    "X": ["Pure", "Standard", "Turbo", "Max"],
-}
-
-
-SUPPORTED_ROBOVAC_MODELS = list(
-    set([item for sublist in ROBOVAC_SERIES.values() for item in sublist])
-)
+from .vacuums import ROBOVAC_MODELS
 
 
 class ModelNotSupportedException(Exception):
@@ -112,48 +11,35 @@ class RoboVac(TuyaDevice):
     """"""
 
     def __init__(self, model_code, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.model_code = model_code
-
-        if self.model_code not in SUPPORTED_ROBOVAC_MODELS:
+        if model_code not in ROBOVAC_MODELS[model_code] is None:
             raise ModelNotSupportedException(
-                "Model {} is not supported".format(self.model_code)
+                "Model {} is not supported".format(model_code)
             )
 
-    def getHomeAssistantFeatures(self):
-        supportedFeatures = (
-            VacuumEntityFeature.BATTERY
-            | VacuumEntityFeature.CLEAN_SPOT
-            | VacuumEntityFeature.FAN_SPEED
-            | VacuumEntityFeature.LOCATE
-            | VacuumEntityFeature.PAUSE
-            | VacuumEntityFeature.RETURN_HOME
-            | VacuumEntityFeature.SEND_COMMAND
-            | VacuumEntityFeature.START
-            | VacuumEntityFeature.STATE
-            | VacuumEntityFeature.STOP
-        )
-
-        if self.model_code in HAS_MAP_FEATURE:
-            supportedFeatures |= VacuumEntityFeature.MAP
+        self.model_details = ROBOVAC_MODELS[model_code]
+        super().__init__(self.model_details, *args, **kwargs)
 
-        return supportedFeatures
+    def getHomeAssistantFeatures(self):
+        return self.model_details.homeassistant_features
 
     def getRoboVacFeatures(self):
-        supportedFeatures = ROBOVAC_SERIES_FEATURES[self.getRoboVacSeries()]
+        return self.model_details.robovac_features
 
-        if self.model_code in HAS_MAP_FEATURE:
-            supportedFeatures |= RoboVacEntityFeature.MAP
+    def getFanSpeeds(self):
+        return self.model_details.commands[RobovacCommand.FAN_SPEED].values
 
-        if self.model_code in HAS_CONSUMABLES:
-            supportedFeatures |= RoboVacEntityFeature.CONSUMABLES
+    def getModes(self):
+        return self.model_details.commands[RobovacCommand.MODE].values
 
-        return supportedFeatures
+    def getSupportedCommands(self):
+        return list(self.model_details.commands.keys())
 
-    def getRoboVacSeries(self):
-        for series, models in ROBOVAC_SERIES.items():
-            if self.model_code in models:
-                return series
+    def getCommandCodes(self):
+        command_codes = {}
+        for key, value in self.model_details.commands:
+            if isinstance(value, dict):
+                command_codes[key] = value.code
+            else:
+                command_codes[key] = value
 
-    def getFanSpeeds(self):
-        return ROBOVAC_SERIES_FAN_SPEEDS[self.getRoboVacSeries()]
+        return command_codes

custom_components/robovac/tuyalocalapi.py

@@ -54,6 +54,8 @@
 from cryptography.hazmat.primitives.hashes import Hash, MD5
 from cryptography.hazmat.primitives.padding import PKCS7
 
+from .vacuums.base import RobovacCommand
+
 INITIAL_BACKOFF = 5
 INITIAL_QUEUE_TIME = 0.1
 BACKOFF_MULTIPLIER = 1.70224
@@ -604,6 +606,7 @@ class TuyaDevice:
 
     def __init__(
         self,
+        model_details,
         device_id,
         host,
         timeout,
@@ -616,6 +619,7 @@ def __init__(
     ):
         """Initialize the device."""
         self._LOGGER = _LOGGER.getChild(device_id)
+        self.model_details = model_details
         self.device_id = device_id
         self.host = host
         self.port = port
@@ -717,7 +721,9 @@ async def async_connect(self):
         try:
             sock.connect((self.host, self.port))
         except (socket.timeout, TimeoutError) as e:
-            self._dps["106"] = "CONNECTION_FAILED"
+            self._dps[self.model_details.commands[RobovacCommand.ERROR]] = (
+                "CONNECTION_FAILED"
+            )
             raise ConnectionTimeoutException("Connection timed out")
         loop = asyncio.get_running_loop()
         loop.create_connection
@@ -744,6 +750,7 @@ async def async_disconnect(self):
 
         if self.writer is not None:
             self.writer.close()
+            await self.writer.wait_closed()
 
         if self.reader is not None and not self.reader.at_eof():
             self.reader.feed_eof()