Validate responses, work around invalid robots (#68)

stianaske · Jan 18, 2021 · b4eda2a · b4eda2a
1 parent 3e555f9
commit b4eda2a
Show file tree

Hide file tree

Showing 4 changed files with 257 additions and 40 deletions.
diff --git a/pybotvac/account.py b/pybotvac/account.py
@@ -5,13 +5,95 @@
 import shutil
 
 import requests
-
-from .exceptions import NeatoRobotException
+from voluptuous import (
+    ALLOW_EXTRA,
+    All,
+    Any,
+    Extra,
+    MultipleInvalid,
+    Optional,
+    Range,
+    Required,
+    Schema,
+    Url,
+)
+
+from .exceptions import NeatoRobotException, NeatoUnsupportedDevice
 from .robot import Robot
 from .session import Session
 
 _LOGGER = logging.getLogger(__name__)
 
+ROBOT_SCHEMA = Schema(
+    {
+        Required("serial"): str,
+        "prefix": Any(str, None),
+        Required("name"): str,
+        "model": str,
+        Required("secret_key"): str,
+        "purchased_at": Any(str, None),
+        "linked_at": Any(str, None),
+        Required("traits"): list,
+        # Everything below this line is not documented, but still present
+        "firmware": str,
+        "timezone": str,
+        Required("nucleo_url"): Url,
+        "mac_address": str,
+        "created_at": str,
+    },
+    extra=ALLOW_EXTRA,
+)
+MAP_SCHEMA = Schema(
+    {
+        "version": int,
+        Required("id"): str,
+        Required("url"): Url,
+        "url_valid_for_seconds": int,
+        Optional("run_id"): str,  # documented, but  not present
+        "status": Any("complete", "incomplete", "canceled"),
+        "launched_from": Any("robot", "schedule", "app"),
+        "error": Any(str, None),
+        "category": int,
+        "mode": int,
+        "modifier": int,
+        "start_at": str,
+        "end_at": str,
+        "end_orientation_relative_degrees": All(int, Range(min=0, max=360)),
+        "run_charge_at_start": All(int, Range(min=0, max=100)),
+        "run_charge_at_end": All(int, Range(min=0, max=100)),
+        "suspended_cleaning_charging_count": int,
+        "time_in_suspended_cleaning": int,
+        "time_in_error": int,
+        "time_in_pause": int,
+        "cleaned_area": float,
+        "base_count": int,
+        "is_docked": bool,
+        "delocalized": bool,
+        # Everything below this line is not documented, but still present
+        "generated_at": str,
+        "persistent_map_id": Any(str, None),
+        "cleaned_with_persistent_map_id": Any(str, None),
+        "valid_as_persistent_map": bool,
+        "navigation_mode": int,
+    },
+    extra=ALLOW_EXTRA,
+)
+MAPS_SCHEMA = Schema(
+    {"stats": {Extra: object}, Required("maps"): [MAP_SCHEMA]},
+    extra=ALLOW_EXTRA,
+)
+PERSISTENT_MAP_SCHEMA = Schema(
+    {
+        Required("id"): str,
+        Required("name"): str,
+        Required("url"): Url,
+        "raw_floor_map_url": Url,
+        "url_valid_for_seconds": int,
+    },
+    extra=ALLOW_EXTRA,
+)
+PERSISTENT_MAPS_SCHEMA = Schema(Required([PERSISTENT_MAP_SCHEMA]))
+
 
 class Account:
     """
@@ -61,8 +143,15 @@ def refresh_maps(self):
         """
 
         for robot in self.robots:
-            resp2 = self._session.get("users/me/robots/{}/maps".format(robot.serial))
-            self._maps.update({robot.serial: resp2.json()})
+            url = f"users/me/robots/{robot.serial}/maps"
+            resp2 = self._session.get(url)
+            try:
+                MAPS_SCHEMA(resp2.json())
+                self._maps.update({robot.serial: resp2.json()})
+            except MultipleInvalid as ex:
+                _LOGGER.warning(
+                    "Invalid response from %s: %s. Got: %s", url, ex, resp2.json()
+                )
 
     def refresh_robots(self):
         """
@@ -74,7 +163,9 @@ def refresh_robots(self):
         resp = self._session.get("users/me/robots")
 
         for robot in resp.json():
+            _LOGGER.debug("Create Robot: %s", robot)
             try:
+                ROBOT_SCHEMA(robot)
                 robot_object = Robot(
                     name=robot["name"],
                     vendor=self._session.vendor,
@@ -84,7 +175,18 @@ def refresh_robots(self):
                     endpoint=robot["nucleo_url"],
                 )
                 self._robots.add(robot_object)
+            except MultipleInvalid as ex:
+                # Robot was not described accordingly by neato
+                _LOGGER.warning(
+                    "Bad response from robots endpoint: %s. Got: %s", ex, robot
+                )
+                continue
+            except NeatoUnsupportedDevice:
+                # Robot does not support home_cleaning service
+                _LOGGER.warning("Your robot %s is unsupported.", robot["name"])
+                continue
             except NeatoRobotException:
+                # The state of the robot could not be received
                 _LOGGER.warning("Your robot %s is offline.", robot["name"])
                 continue
 
@@ -143,8 +245,13 @@ def refresh_persistent_maps(self):
         """
 
         for robot in self._robots:
-            resp2 = self._session.get(
-                "users/me/robots/{}/persistent_maps".format(robot.serial)
-            )
+            url = f"users/me/robots/{robot.serial}/persistent_maps"
+            resp2 = self._session.get(url)
 
-            self._persistent_maps.update({robot.serial: resp2.json()})
+            try:
+                PERSISTENT_MAPS_SCHEMA(resp2.json())
+                self._persistent_maps.update({robot.serial: resp2.json()})
+            except MultipleInvalid as ex:
+                _LOGGER.warning(
+                    "Invalid response from %s: %s. Got: %s", url, ex, resp2.json()
+                )