Fix Neato reauth flow when token expired (#52843)

* Fix Neato reauth flow when token expired

* Change and simplify approach

* Missing file

* Cleanup

* Update unique_id

* Added missing lamda

* Unique_id reworked

* Guard for id future ID changes

* Bump pybotvac: provide unique_id

* Address review comment

* Fix update check

* Remove token check

* Trigger reauth only for 401 and 403 code response

* Review comments

.coveragerc

@@ -677,6 +677,7 @@ omit =
     homeassistant/components/neato/__init__.py
     homeassistant/components/neato/api.py
     homeassistant/components/neato/camera.py
+    homeassistant/components/neato/hub.py
     homeassistant/components/neato/sensor.py
     homeassistant/components/neato/switch.py
     homeassistant/components/neato/vacuum.py

homeassistant/components/neato/__init__.py

@@ -1,7 +1,7 @@
 """Support for Neato botvac connected vacuum cleaners."""
-from datetime import timedelta
 import logging
 
+import aiohttp
 from pybotvac import Account, Neato
 from pybotvac.exceptions import NeatoException
 import voluptuous as vol
@@ -12,17 +12,10 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers import config_entry_oauth2_flow, config_validation as cv
 from homeassistant.helpers.typing import ConfigType
-from homeassistant.util import Throttle
 
 from . import api, config_flow
-from .const import (
+from .const import NEATO_CONFIG, NEATO_DOMAIN, NEATO_LOGIN
-    NEATO_CONFIG,
+from .hub import NeatoHub
-    NEATO_DOMAIN,
-    NEATO_LOGIN,
-    NEATO_MAP_DATA,
-    NEATO_PERSISTENT_MAPS,
-    NEATO_ROBOTS,
-)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -77,10 +70,20 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         )
     )
 
+    session = config_entry_oauth2_flow.OAuth2Session(hass, entry, implementation)
+    try:
+        await session.async_ensure_token_valid()
+    except aiohttp.ClientResponseError as ex:
+        _LOGGER.debug("API error: %s (%s)", ex.code, ex.message)
+        if ex.code in (401, 403):
+            raise ConfigEntryAuthFailed("Token not valid, trigger renewal") from ex
+
     neato_session = api.ConfigEntryAuth(hass, entry, implementation)
     hass.data[NEATO_DOMAIN][entry.entry_id] = neato_session
     hub = NeatoHub(hass, Account(neato_session))
 
+    await hub.async_update_entry_unique_id(entry)
+
     try:
         await hass.async_add_executor_job(hub.update_robots)
     except NeatoException as ex:
@@ -94,32 +97,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigType) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload config entry."""
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
     if unload_ok:
         hass.data[NEATO_DOMAIN].pop(entry.entry_id)
 
     return unload_ok
-
-
-class NeatoHub:
-    """A My Neato hub wrapper class."""
-
-    def __init__(self, hass: HomeAssistant, neato: Account) -> None:
-        """Initialize the Neato hub."""
-        self._hass = hass
-        self.my_neato: Account = neato
-
-    @Throttle(timedelta(minutes=1))
-    def update_robots(self):
-        """Update the robot states."""
-        _LOGGER.debug("Running HUB.update_robots %s", self._hass.data.get(NEATO_ROBOTS))
-        self._hass.data[NEATO_ROBOTS] = self.my_neato.robots
-        self._hass.data[NEATO_PERSISTENT_MAPS] = self.my_neato.persistent_maps
-        self._hass.data[NEATO_MAP_DATA] = self.my_neato.maps
-
-    def download_map(self, url):
-        """Download a new map image."""
-        map_image_data = self.my_neato.get_map_image(url)
-        return map_image_data

homeassistant/components/neato/config_flow.py

@@ -5,7 +5,7 @@ import logging
 
 import voluptuous as vol
 
-from homeassistant.const import CONF_TOKEN
+from homeassistant.config_entries import SOURCE_REAUTH
 from homeassistant.helpers import config_entry_oauth2_flow
 
 from .const import NEATO_DOMAIN
@@ -26,7 +26,7 @@ class OAuth2FlowHandler(
     async def async_step_user(self, user_input: dict | None = None) -> dict:
         """Create an entry for the flow."""
         current_entries = self._async_current_entries()
-        if current_entries and CONF_TOKEN in current_entries[0].data:
+        if self.source != SOURCE_REAUTH and current_entries:
             # Already configured
             return self.async_abort(reason="already_configured")
 
@@ -47,7 +47,7 @@ class OAuth2FlowHandler(
     async def async_oauth_create_entry(self, data: dict) -> dict:
         """Create an entry for the flow. Update an entry if one already exist."""
         current_entries = self._async_current_entries()
-        if current_entries and CONF_TOKEN not in current_entries[0].data:
+        if self.source == SOURCE_REAUTH and current_entries:
             # Update entry
             self.hass.config_entries.async_update_entry(
                 current_entries[0], title=self.flow_impl.name, data=data

homeassistant/components/neato/hub.py

@@ -0,0 +1,47 @@
+"""Support for Neato botvac connected vacuum cleaners."""
+from datetime import timedelta
+import logging
+
+from pybotvac import Account
+
+from homeassistant.core import HomeAssistant
+from homeassistant.util import Throttle
+
+from .const import NEATO_MAP_DATA, NEATO_PERSISTENT_MAPS, NEATO_ROBOTS
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class NeatoHub:
+    """A My Neato hub wrapper class."""
+
+    def __init__(self, hass: HomeAssistant, neato: Account) -> None:
+        """Initialize the Neato hub."""
+        self._hass = hass
+        self.my_neato: Account = neato
+
+    @Throttle(timedelta(minutes=1))
+    def update_robots(self):
+        """Update the robot states."""
+        _LOGGER.debug("Running HUB.update_robots %s", self._hass.data.get(NEATO_ROBOTS))
+        self._hass.data[NEATO_ROBOTS] = self.my_neato.robots
+        self._hass.data[NEATO_PERSISTENT_MAPS] = self.my_neato.persistent_maps
+        self._hass.data[NEATO_MAP_DATA] = self.my_neato.maps
+
+    def download_map(self, url):
+        """Download a new map image."""
+        map_image_data = self.my_neato.get_map_image(url)
+        return map_image_data
+
+    async def async_update_entry_unique_id(self, entry) -> str:
+        """Update entry for unique_id."""
+
+        await self._hass.async_add_executor_job(self.my_neato.refresh_userdata)
+        unique_id = self.my_neato.unique_id
+
+        if entry.unique_id == unique_id:
+            return unique_id
+
+        _LOGGER.debug("Updating user unique_id for previous config entry")
+        self._hass.config_entries.async_update_entry(entry, unique_id=unique_id)
+        return unique_id

homeassistant/components/neato/manifest.json

@@ -3,7 +3,7 @@
   "name": "Neato Botvac",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/neato",
-  "requirements": ["pybotvac==0.0.21"],
+  "requirements": ["pybotvac==0.0.22"],
   "codeowners": ["@dshokouhi", "@Santobert"],
   "dependencies": ["http"],
   "iot_class": "cloud_polling"

requirements_all.txt

@@ -1342,7 +1342,7 @@ pyblackbird==0.5
 # pybluez==0.22
 
 # homeassistant.components.neato
-pybotvac==0.0.21
+pybotvac==0.0.22
 
 # homeassistant.components.nissan_leaf
 pycarwings2==2.10

requirements_test_all.txt

@@ -760,7 +760,7 @@ pyatv==0.8.2
 pyblackbird==0.5
 
 # homeassistant.components.neato
-pybotvac==0.0.21
+pybotvac==0.0.22
 
 # homeassistant.components.cloudflare
 pycfdns==1.2.1