From ddfbeffd2857d183d589d6f09510918725af482a Mon Sep 17 00:00:00 2001
From: Daniel Shokouhi <dshokouhi@gmail.com>
Date: Tue, 14 Jul 2020 13:59:03 -0700
Subject: [PATCH] Fix zone cleaning and raise config entry not ready when
 needed (#37741)

---
 homeassistant/components/neato/__init__.py | 11 +++++----
 homeassistant/components/neato/const.py    |  2 +-
 homeassistant/components/neato/vacuum.py   | 26 +++++++++++++++++-----
 3 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/homeassistant/components/neato/__init__.py b/homeassistant/components/neato/__init__.py
index ad4eb02eccc..a80d555708c 100644
--- a/homeassistant/components/neato/__init__.py
+++ b/homeassistant/components/neato/__init__.py
@@ -9,6 +9,7 @@ import voluptuous as vol
 
 from homeassistant.config_entries import SOURCE_IMPORT
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
+from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import config_validation as cv
 from homeassistant.util import Throttle
 
@@ -91,9 +92,8 @@ async def async_setup(hass, config):
 
 async def async_setup_entry(hass, entry):
     """Set up config entry."""
-    hass.data[NEATO_LOGIN] = NeatoHub(hass, entry.data, Account)
+    hub = NeatoHub(hass, entry.data, Account)
 
-    hub = hass.data[NEATO_LOGIN]
     await hass.async_add_executor_job(hub.login)
     if not hub.logged_in:
         _LOGGER.debug("Failed to login to Neato API")
@@ -103,10 +103,12 @@ async def async_setup_entry(hass, entry):
         await hass.async_add_executor_job(hub.update_robots)
     except NeatoRobotException:
         _LOGGER.debug("Failed to connect to Neato API")
-        return False
+        raise ConfigEntryNotReady
+
+    hass.data[NEATO_LOGIN] = hub
 
     for component in ("camera", "vacuum", "switch", "sensor"):
-        hass.async_create_task(
+        hass.async_add_job(
             hass.config_entries.async_forward_entry_setup(entry, component)
         )
 
@@ -154,6 +156,7 @@ class NeatoHub:
                 _LOGGER.error("Invalid credentials")
             else:
                 _LOGGER.error("Unable to connect to Neato API")
+                raise ConfigEntryNotReady
             self.logged_in = False
             return
 
diff --git a/homeassistant/components/neato/const.py b/homeassistant/components/neato/const.py
index cfe8a2dad9d..144ea40b92a 100644
--- a/homeassistant/components/neato/const.py
+++ b/homeassistant/components/neato/const.py
@@ -9,7 +9,7 @@ NEATO_MAP_DATA = "neato_map_data"
 NEATO_PERSISTENT_MAPS = "neato_persistent_maps"
 NEATO_ROBOTS = "neato_robots"
 
-SCAN_INTERVAL_MINUTES = 5
+SCAN_INTERVAL_MINUTES = 1
 
 SERVICE_NEATO_CUSTOM_CLEANING = "custom_cleaning"
 
diff --git a/homeassistant/components/neato/vacuum.py b/homeassistant/components/neato/vacuum.py
index a67b48169c4..90b52aab073 100644
--- a/homeassistant/components/neato/vacuum.py
+++ b/homeassistant/components/neato/vacuum.py
@@ -152,7 +152,7 @@ class NeatoConnectedVacuum(StateVacuumEntity):
         self._clean_error_time = None
         self._launched_from = None
         self._battery_level = None
-        self._robot_boundaries = {}
+        self._robot_boundaries = []
         self._robot_stats = None
 
     def update(self):
@@ -241,14 +241,27 @@ class NeatoConnectedVacuum(StateVacuumEntity):
             and self._robot_maps[self._robot_serial]
         ):
             allmaps = self._robot_maps[self._robot_serial]
+            _LOGGER.debug("Found the following maps for '%s': %s", self._name, allmaps)
+            self._robot_boundaries = []  # Reset boundaries before refreshing boundaries
             for maps in allmaps:
                 try:
-                    self._robot_boundaries = self.robot.get_map_boundaries(
-                        maps["id"]
-                    ).json()
+                    robot_boundaries = self.robot.get_map_boundaries(maps["id"]).json()
                 except NeatoRobotException as ex:
                     _LOGGER.error("Could not fetch map boundaries: %s", ex)
-                    self._robot_boundaries = {}
+                    return
+
+                _LOGGER.debug(
+                    "Boundaries for robot '%s' in map '%s': %s",
+                    self._name,
+                    maps["name"],
+                    robot_boundaries,
+                )
+                self._robot_boundaries += robot_boundaries["data"]["boundaries"]
+                _LOGGER.debug(
+                    "List of boundaries for '%s': %s",
+                    self._name,
+                    self._robot_boundaries,
+                )
 
     @property
     def name(self):
@@ -323,6 +336,7 @@ class NeatoConnectedVacuum(StateVacuumEntity):
             info["manufacturer"] = self._robot_stats["battery"]["vendor"]
             info["model"] = self._robot_stats["model"]
             info["sw_version"] = self._robot_stats["firmware"]
+        return info
 
     def start(self):
         """Start cleaning or resume cleaning."""
@@ -376,7 +390,7 @@ class NeatoConnectedVacuum(StateVacuumEntity):
         """Zone cleaning service call."""
         boundary_id = None
         if zone is not None:
-            for boundary in self._robot_boundaries["data"]["boundaries"]:
+            for boundary in self._robot_boundaries:
                 if zone in boundary["name"]:
                     boundary_id = boundary["id"]
             if boundary_id is None: