Migrate neato to use runtime_data (#166854)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

homeassistant/components/neato/__init__.py

@@ -24,12 +24,14 @@ from homeassistant.helpers.config_entry_oauth2_flow import (
 from homeassistant.helpers.typing import ConfigType
 
 from . import api
-from .const import DOMAIN, NEATO_LOGIN
+from .const import DOMAIN
 from .hub import NeatoHub
 from .services import async_setup_services
 
 _LOGGER = logging.getLogger(__name__)
 
+type NeatoConfigEntry = ConfigEntry[NeatoHub]
+
 CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
 PLATFORMS = [
     Platform.BUTTON,
@@ -46,9 +48,8 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: NeatoConfigEntry) -> bool:
     """Set up config entry."""
-    hass.data.setdefault(DOMAIN, {})
     if CONF_TOKEN not in entry.data:
         raise ConfigEntryAuthFailed
 
@@ -69,7 +70,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         raise ConfigEntryNotReady from ex
 
     neato_session = api.ConfigEntryAuth(hass, entry, implementation)
-    hass.data[DOMAIN][entry.entry_id] = neato_session
     hub = NeatoHub(hass, Account(neato_session))
 
     await hub.async_update_entry_unique_id(entry)
@@ -80,17 +80,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         _LOGGER.debug("Failed to connect to Neato API")
         raise ConfigEntryNotReady from ex
 
-    hass.data[NEATO_LOGIN] = hub
+    entry.runtime_data = hub
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: NeatoConfigEntry) -> bool:
     """Unload config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/neato/button.py

@@ -5,22 +5,21 @@ from __future__ import annotations
 from pybotvac import Robot
 
 from homeassistant.components.button import ButtonEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import NEATO_ROBOTS
+from . import NeatoConfigEntry
 from .entity import NeatoEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: NeatoConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Neato button from config entry."""
-    entities = [NeatoDismissAlertButton(robot) for robot in hass.data[NEATO_ROBOTS]]
+    entities = [NeatoDismissAlertButton(robot) for robot in entry.runtime_data.robots]
 
     async_add_entities(entities, True)
 

homeassistant/components/neato/camera.py

@@ -11,11 +11,11 @@ from pybotvac.robot import Robot
 from urllib3.response import HTTPResponse
 
 from homeassistant.components.camera import Camera
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import NEATO_LOGIN, NEATO_MAP_DATA, NEATO_ROBOTS, SCAN_INTERVAL_MINUTES
+from . import NeatoConfigEntry
+from .const import SCAN_INTERVAL_MINUTES
 from .entity import NeatoEntity
 from .hub import NeatoHub
 
@@ -27,15 +27,14 @@ ATTR_GENERATED_AT = "generated_at"
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: NeatoConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Neato camera with config entry."""
-    neato: NeatoHub = hass.data[NEATO_LOGIN]
-    mapdata: dict[str, Any] | None = hass.data.get(NEATO_MAP_DATA)
+    hub = entry.runtime_data
     dev = [
-        NeatoCleaningMap(neato, robot, mapdata)
-        for robot in hass.data[NEATO_ROBOTS]
+        NeatoCleaningMap(hub, robot, hub.map_data)
+        for robot in hub.robots
         if "maps" in robot.traits
     ]
 
@@ -51,9 +50,7 @@ class NeatoCleaningMap(NeatoEntity, Camera):
 
     _attr_translation_key = "cleaning_map"
 
-    def __init__(
-        self, neato: NeatoHub, robot: Robot, mapdata: dict[str, Any] | None
-    ) -> None:
+    def __init__(self, neato: NeatoHub, robot: Robot, mapdata: dict[str, Any]) -> None:
         """Initialize Neato cleaning map."""
         super().__init__(robot)
         Camera.__init__(self)

homeassistant/components/neato/const.py

@@ -3,10 +3,6 @@
 DOMAIN = "neato"
 
 CONF_VENDOR = "vendor"
-NEATO_LOGIN = "neato_login"
-NEATO_MAP_DATA = "neato_map_data"
-NEATO_PERSISTENT_MAPS = "neato_persistent_maps"
-NEATO_ROBOTS = "neato_robots"
 
 SCAN_INTERVAL_MINUTES = 1
 

homeassistant/components/neato/hub.py

@@ -1,7 +1,10 @@
 """Support for Neato botvac connected vacuum cleaners."""
 
+from __future__ import annotations
+
 from datetime import timedelta
 import logging
+from typing import Any
 
 from pybotvac import Account
 from urllib3.response import HTTPResponse
@@ -10,8 +13,6 @@ from homeassistant.config_entries import ConfigEntry
 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__)
 
 
@@ -22,14 +23,17 @@ class NeatoHub:
         """Initialize the Neato hub."""
         self._hass = hass
         self.my_neato: Account = neato
+        self.robots: set[Any] = set()
+        self.persistent_maps: dict[str, Any] = {}
+        self.map_data: dict[str, Any] = {}
 
     @Throttle(timedelta(minutes=1))
     def update_robots(self) -> None:
         """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
+        _LOGGER.debug("Running HUB.update_robots %s", self.robots)
+        self.robots = self.my_neato.robots
+        self.persistent_maps = self.my_neato.persistent_maps
+        self.map_data = self.my_neato.maps
 
     def download_map(self, url: str) -> HTTPResponse:
         """Download a new map image."""

homeassistant/components/neato/sensor.py

@@ -10,12 +10,12 @@ from pybotvac.exceptions import NeatoRobotException
 from pybotvac.robot import Robot
 
 from homeassistant.components.sensor import SensorDeviceClass, SensorEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE, EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import NEATO_LOGIN, NEATO_ROBOTS, SCAN_INTERVAL_MINUTES
+from . import NeatoConfigEntry
+from .const import SCAN_INTERVAL_MINUTES
 from .entity import NeatoEntity
 from .hub import NeatoHub
 
@@ -28,12 +28,12 @@ BATTERY = "Battery"
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: NeatoConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the Neato sensor using config entry."""
-    neato: NeatoHub = hass.data[NEATO_LOGIN]
-    dev = [NeatoSensor(neato, robot) for robot in hass.data[NEATO_ROBOTS]]
+    hub = entry.runtime_data
+    dev = [NeatoSensor(hub, robot) for robot in hub.robots]
 
     if not dev:
         return

homeassistant/components/neato/switch.py

@@ -10,12 +10,12 @@ from pybotvac.exceptions import NeatoRobotException
 from pybotvac.robot import Robot
 
 from homeassistant.components.switch import SwitchEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import STATE_OFF, STATE_ON, EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import NEATO_LOGIN, NEATO_ROBOTS, SCAN_INTERVAL_MINUTES
+from . import NeatoConfigEntry
+from .const import SCAN_INTERVAL_MINUTES
 from .entity import NeatoEntity
 from .hub import NeatoHub
 
@@ -30,14 +30,14 @@ SWITCH_TYPES = {SWITCH_TYPE_SCHEDULE: ["Schedule"]}
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: NeatoConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Neato switch with config entry."""
-    neato: NeatoHub = hass.data[NEATO_LOGIN]
+    hub = entry.runtime_data
     dev = [
-        NeatoConnectedSwitch(neato, robot, type_name)
-        for robot in hass.data[NEATO_ROBOTS]
+        NeatoConnectedSwitch(hub, robot, type_name)
+        for robot in hub.robots
         for type_name in SWITCH_TYPES
     ]
 

homeassistant/components/neato/vacuum.py

@@ -15,22 +15,12 @@ from homeassistant.components.vacuum import (
     VacuumActivity,
     VacuumEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import (
-    ACTION,
-    ALERTS,
-    ERRORS,
-    MODE,
-    NEATO_LOGIN,
-    NEATO_MAP_DATA,
-    NEATO_PERSISTENT_MAPS,
-    NEATO_ROBOTS,
-    SCAN_INTERVAL_MINUTES,
-)
+from . import NeatoConfigEntry
+from .const import ACTION, ALERTS, ERRORS, MODE, SCAN_INTERVAL_MINUTES
 from .entity import NeatoEntity
 from .hub import NeatoHub
 
@@ -52,16 +42,16 @@ ATTR_LAUNCHED_FROM = "launched_from"
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: NeatoConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Neato vacuum with config entry."""
-    neato: NeatoHub = hass.data[NEATO_LOGIN]
-    mapdata: dict[str, Any] | None = hass.data.get(NEATO_MAP_DATA)
-    persistent_maps: dict[str, Any] | None = hass.data.get(NEATO_PERSISTENT_MAPS)
+    hub = entry.runtime_data
     dev = [
-        NeatoConnectedVacuum(neato, robot, mapdata, persistent_maps)
-        for robot in hass.data[NEATO_ROBOTS]
+        NeatoConnectedVacuum(
+            hub, robot, hub.map_data or None, hub.persistent_maps or None
+        )
+        for robot in hub.robots
     ]
 
     if not dev: