From 711997cb13529a78e332757a92f65fad55f323b9 Mon Sep 17 00:00:00 2001
From: Luke Bonaccorsi <luke@lukeb.co.uk>
Date: Fri, 18 Aug 2023 12:13:01 +0100
Subject: [PATCH] Add errors for no IP and no connection

---
 custom_components/robovac/errors.py       | 63 ++++++++++++-----------
 custom_components/robovac/tuyalocalapi.py |  2 +
 custom_components/robovac/vacuum.py       |  4 +-
 3 files changed, 37 insertions(+), 32 deletions(-)

diff --git a/custom_components/robovac/errors.py b/custom_components/robovac/errors.py
index 8a5fabd..c79181f 100644
--- a/custom_components/robovac/errors.py
+++ b/custom_components/robovac/errors.py
@@ -1,37 +1,38 @@
 ERROR_MESSAGES = {
     "IP_ADDRESS": "IP Address not set",
+    "CONNECTION_FAILED": "Connection to the vacuum failed",
     "no_error": "None",
-    1:"Error: Front bumper stuck",
-    2:"Error: Wheel stuck",
-    3:"Error: Side brush",
-    4:"Error: Rolling brush bar stuck",
-    5:"Error: Device trapped",
-    6:"Error: Device trapped",
-    7:"Error: Wheel suspended",
-    8:"Error: Low battery",
-    9:"Error: Magnetic boundary",
-    12:"Error: Right wall sensor",
-    13:"Error: Device tilted",
-    14:"Error: Insert dust collector",
-    17:"Error: Restricted area detected",
-    18:"Error: Laser cover stuck",
-    19:"Error: Laser sesor stuck",
-    20:"Error: Laser sensor blocked",
-    21:"Error: Base blocked",
-    "S1":"Error: Battery",
-    "S2":"Error: Wheel Module",
-    "S3":"Error: Side Brush",
-    "S4":"Error: Suction Fan",
-    "S5":"Error: Rolling Brush",
-    "S8":"Error: Path Tracking Sensor",
-    "Wheel_stuck":"Error: Wheel stuck",
-    "R_brush_stuck":"Error: Rolling brush stuck",
-    "Crash_bar_stuck":"Error: Front bumper stuck",
-    "sensor_dirty":"Error: Sensor dirty",
-    "N_enough_pow":"Error: Low battery",
-    "Stuck_5_min":"Error: Device trapped",
-    "Fan_stuck":"Error: Fan stuck",
-    "S_brush_stuck":"Error: Side brush stuck",
+    1:"Front bumper stuck",
+    2:"Wheel stuck",
+    3:"Side brush",
+    4:"Rolling brush bar stuck",
+    5:"Device trapped",
+    6:"Device trapped",
+    7:"Wheel suspended",
+    8:"Low battery",
+    9:"Magnetic boundary",
+    12:"Right wall sensor",
+    13:"Device tilted",
+    14:"Insert dust collector",
+    17:"Restricted area detected",
+    18:"Laser cover stuck",
+    19:"Laser sesor stuck",
+    20:"Laser sensor blocked",
+    21:"Base blocked",
+    "S1":"Battery",
+    "S2":"Wheel Module",
+    "S3":"Side Brush",
+    "S4":"Suction Fan",
+    "S5":"Rolling Brush",
+    "S8":"Path Tracking Sensor",
+    "Wheel_stuck":"Wheel stuck",
+    "R_brush_stuck":"Rolling brush stuck",
+    "Crash_bar_stuck":"Front bumper stuck",
+    "sensor_dirty":"Sensor dirty",
+    "N_enough_pow":"Low battery",
+    "Stuck_5_min":"Device trapped",
+    "Fan_stuck":"Fan stuck",
+    "S_brush_stuck":"Side brush stuck",
 }
 
 def getErrorMessage(code):
diff --git a/custom_components/robovac/tuyalocalapi.py b/custom_components/robovac/tuyalocalapi.py
index c303ba7..2521a21 100644
--- a/custom_components/robovac/tuyalocalapi.py
+++ b/custom_components/robovac/tuyalocalapi.py
@@ -677,6 +677,7 @@ class TuyaDevice:
         try:
             sock.connect((self.host, self.port))
         except socket.timeout as e:
+            self._dps["106"] = "CONNECTION_FAILED"
             raise ConnectionTimeoutException("Connection timed out") from e
         self.reader, self.writer = await asyncio.open_connection(sock=sock)
         self._connected = True
@@ -740,6 +741,7 @@ class TuyaDevice:
             response_data = await self.reader.readuntil(MAGIC_SUFFIX_BYTES)
         except socket.error as e:
             _LOGGER.error("Connection to {} failed: {}".format(self, e))
+            self._dps["106"] = "CONNECTION_FAILED"
             asyncio.ensure_future(self.async_disconnect())
             return
         except asyncio.IncompleteReadError as e:
diff --git a/custom_components/robovac/vacuum.py b/custom_components/robovac/vacuum.py
index 05a07f4..dca1230 100644
--- a/custom_components/robovac/vacuum.py
+++ b/custom_components/robovac/vacuum.py
@@ -201,7 +201,8 @@ class RoboVacEntity(StateVacuumEntity):
     def extra_state_attributes(self) -> dict[str, Any]:
         """Return the device-specific state attributes of this vacuum."""
         data: dict[str, Any] = {}
-        data[ATTR_ERROR] = getErrorMessage(self.error_code)
+        if type(self.error_code) is not None and self.error_code not in [0, "no_error"]:
+            data[ATTR_ERROR] = getErrorMessage(self.error_code)
 
         if self.supported_features & VacuumEntityFeature.STATUS:
             data[ATTR_STATUS] = self.status
@@ -263,6 +264,7 @@ class RoboVacEntity(StateVacuumEntity):
         """Synchronise state from the vacuum."""
         self.async_write_ha_state()
         if self.ip_address == "":
+            self.error_code = "IP_ADDRESS"
             return
         await self.vacuum.async_get()
         self.tuyastatus = self.vacuum._dps