Store runtime data in entry in Ecovacs

homeassistant/components/ecovacs/__init__.py

@@ -37,6 +37,7 @@
  Platform.SWITCH,
  Platform.VACUUM,
 ]
+EcovacsConfigEntry = ConfigEntry[EcovacsController]
 
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
@@ -50,21 +51,20 @@ 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: EcovacsConfigEntry) -> bool:
  """Set up this integration using UI."""
  controller = EcovacsController(hass, entry.data)
  await controller.initialize()
 
- hass.data.setdefault(DOMAIN, {})[entry.entry_id] = controller
+ async def on_unload() -> None:
+ await controller.teardown()
+
+ entry.async_on_unload(on_unload)
+ entry.runtime_data = controller
  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: EcovacsConfigEntry) -> bool:
  """Unload config entry."""
- if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
- await hass.data[DOMAIN][entry.entry_id].teardown()
- hass.data[DOMAIN].pop(entry.entry_id)
- if not hass.data[DOMAIN]:
- hass.data.pop(DOMAIN)
- return unload_ok
+ return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/ecovacs/binary_sensor.py

@@ -11,13 +11,11 @@
  BinarySensorEntity,
  BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .controller import EcovacsController
+from . import EcovacsConfigEntry
 from .entity import (
  CapabilityDevice,
  EcovacsCapabilityEntityDescription,
@@ -52,13 +50,14 @@ class EcovacsBinarySensorEntityDescription(
 
 async def async_setup_entry(
  hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: EcovacsConfigEntry,
  async_add_entities: AddEntitiesCallback,
 ) -> None:
  """Add entities for passed config_entry in HA."""
- controller: EcovacsController = hass.data[DOMAIN][config_entry.entry_id]
  async_add_entities(
- get_supported_entitites(controller, EcovacsBinarySensor, ENTITY_DESCRIPTIONS)
+ get_supported_entitites(
+ config_entry.runtime_data, EcovacsBinarySensor, ENTITY_DESCRIPTIONS
+ )
  )
 
 

homeassistant/components/ecovacs/button.py

@@ -11,13 +11,12 @@
 from deebot_client.events import LifeSpan
 
 from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN, SUPPORTED_LIFESPANS
-from .controller import EcovacsController
+from . import EcovacsConfigEntry
+from .const import SUPPORTED_LIFESPANS
 from .entity import (
  CapabilityDevice,
  EcovacsCapabilityEntityDescription,
@@ -66,11 +65,11 @@ class EcovacsLifespanButtonEntityDescription(ButtonEntityDescription):
 
 async def async_setup_entry(
  hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: EcovacsConfigEntry,
  async_add_entities: AddEntitiesCallback,
 ) -> None:
  """Add entities for passed config_entry in HA."""
- controller: EcovacsController = hass.data[DOMAIN][config_entry.entry_id]
+ controller = config_entry.runtime_data
  entities: list[EcovacsEntity] = get_supported_entitites(
  controller, EcovacsButtonEntity, ENTITY_DESCRIPTIONS
  )

homeassistant/components/ecovacs/controller.py

@@ -42,7 +42,7 @@ def __init__(self, hass: HomeAssistant, config: Mapping[str, Any]) -> None:
  """Initialize controller."""
  self._hass = hass
  self._devices: list[Device] = []
- self.legacy_devices: list[VacBot] = []
+ self._legacy_devices: list[VacBot] = []
  rest_url = config.get(CONF_OVERRIDE_REST_URL)
  self._device_id = get_client_device_id(hass, rest_url is not None)
  country = config[CONF_COUNTRY]
@@ -101,7 +101,7 @@ async def initialize(self) -> None:
  self._continent,
  monitor=True,
  )
- self.legacy_devices.append(bot)
+ self._legacy_devices.append(bot)
  except InvalidAuthenticationError as ex:
  raise ConfigEntryError("Invalid credentials") from ex
  except DeebotError as ex:
@@ -113,7 +113,7 @@ async def teardown(self) -> None:
  """Disconnect controller."""
  for device in self._devices:
  await device.teardown()
- for legacy_device in self.legacy_devices:
+ for legacy_device in self._legacy_devices:
  await self._hass.async_add_executor_job(legacy_device.disconnect)
  await self._mqtt.disconnect()
  await self._authenticator.teardown()
@@ -124,3 +124,8 @@ def devices(self, capability: type[Capabilities]) -> Generator[Device, None, Non
  for device in self._devices:
  if isinstance(device.capabilities, capability):
  yield device
+
+ @property
+ def legacy_devices(self) -> list[VacBot]:
+ """Return legacy devices."""
+ return self._legacy_devices

homeassistant/components/ecovacs/diagnostics.py

@@ -7,12 +7,11 @@
 from deebot_client.capabilities import Capabilities
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME, CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 
-from .const import CONF_OVERRIDE_MQTT_URL, CONF_OVERRIDE_REST_URL, DOMAIN
-from .controller import EcovacsController
+from . import EcovacsConfigEntry
+from .const import CONF_OVERRIDE_MQTT_URL, CONF_OVERRIDE_REST_URL
 
 REDACT_CONFIG = {
  CONF_USERNAME,
@@ -25,10 +24,10 @@
 
 
 async def async_get_config_entry_diagnostics(
- hass: HomeAssistant, config_entry: ConfigEntry
+ hass: HomeAssistant, config_entry: EcovacsConfigEntry
 ) -> dict[str, Any]:
  """Return diagnostics for a config entry."""
- controller: EcovacsController = hass.data[DOMAIN][config_entry.entry_id]
+ controller = config_entry.runtime_data
  diag: dict[str, Any] = {
  "config": async_redact_data(config_entry.as_dict(), REDACT_CONFIG)
  }

homeassistant/components/ecovacs/event.py

@@ -5,24 +5,22 @@
 from deebot_client.events import CleanJobStatus, ReportStatsEvent
 
 from homeassistant.components.event import EventEntity, EventEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .controller import EcovacsController
+from . import EcovacsConfigEntry
 from .entity import EcovacsEntity
 from .util import get_name_key
 
 
 async def async_setup_entry(
  hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: EcovacsConfigEntry,
  async_add_entities: AddEntitiesCallback,
 ) -> None:
  """Add entities for passed config_entry in HA."""
- controller: EcovacsController = hass.data[DOMAIN][config_entry.entry_id]
+ controller = config_entry.runtime_data
  async_add_entities(
  EcovacsLastJobEventEntity(device) for device in controller.devices(Capabilities)
  )

homeassistant/components/ecovacs/image.py

@@ -5,23 +5,21 @@
 from deebot_client.events.map import CachedMapInfoEvent, MapChangedEvent
 
 from homeassistant.components.image import ImageEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity import EntityDescription
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .controller import EcovacsController
+from . import EcovacsConfigEntry
 from .entity import EcovacsEntity
 
 
 async def async_setup_entry(
  hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: EcovacsConfigEntry,
  async_add_entities: AddEntitiesCallback,
 ) -> None:
  """Add entities for passed config_entry in HA."""
- controller: EcovacsController = hass.data[DOMAIN][config_entry.entry_id]
+ controller = config_entry.runtime_data
  entities = []
  for device in controller.devices(VacuumCapabilities):
  capabilities: VacuumCapabilities = device.capabilities

homeassistant/components/ecovacs/lawn_mower.py

@@ -15,12 +15,10 @@
  LawnMowerEntityEntityDescription,
  LawnMowerEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .controller import EcovacsController
+from . import EcovacsConfigEntry
 from .entity import EcovacsEntity
 
 _LOGGER = logging.getLogger(__name__)
@@ -38,11 +36,11 @@
 
 async def async_setup_entry(
  hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: EcovacsConfigEntry,
  async_add_entities: AddEntitiesCallback,
 ) -> None:
  """Set up the Ecovacs mowers."""
- controller: EcovacsController = hass.data[DOMAIN][config_entry.entry_id]
+ controller = config_entry.runtime_data
  mowers: list[EcovacsMower] = [
  EcovacsMower(device) for device in controller.devices(MowerCapabilities)
  ]

homeassistant/components/ecovacs/number.py

@@ -10,13 +10,11 @@
 from deebot_client.events import CleanCountEvent, VolumeEvent
 
 from homeassistant.components.number import NumberEntity, NumberEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .controller import EcovacsController
+from . import EcovacsConfigEntry
 from .entity import (
  CapabilityDevice,
  EcovacsCapabilityEntityDescription,
@@ -70,11 +68,11 @@ class EcovacsNumberEntityDescription(
 
 async def async_setup_entry(
  hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: EcovacsConfigEntry,
  async_add_entities: AddEntitiesCallback,
 ) -> None:
  """Add entities for passed config_entry in HA."""
- controller: EcovacsController = hass.data[DOMAIN][config_entry.entry_id]
+ controller = config_entry.runtime_data
  entities: list[EcovacsEntity] = get_supported_entitites(
  controller, EcovacsNumberEntity, ENTITY_DESCRIPTIONS
  )

homeassistant/components/ecovacs/select.py

@@ -9,13 +9,11 @@
 from deebot_client.events import WaterInfoEvent, WorkModeEvent
 
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .controller import EcovacsController
+from . import EcovacsConfigEntry
 from .entity import (
  CapabilityDevice,
  EcovacsCapabilityEntityDescription,
@@ -62,11 +60,11 @@ class EcovacsSelectEntityDescription(
 
 async def async_setup_entry(
  hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: EcovacsConfigEntry,
  async_add_entities: AddEntitiesCallback,
 ) -> None:
  """Add entities for passed config_entry in HA."""
- controller: EcovacsController = hass.data[DOMAIN][config_entry.entry_id]
+ controller = config_entry.runtime_data
  entities = get_supported_entitites(
  controller, EcovacsSelectEntity, ENTITY_DESCRIPTIONS
  )

homeassistant/components/ecovacs/sensor.py

@@ -24,7 +24,6 @@
  SensorEntityDescription,
  SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
  AREA_SQUARE_METERS,
  ATTR_BATTERY_LEVEL,
@@ -37,8 +36,8 @@
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
-from .const import DOMAIN, SUPPORTED_LIFESPANS
-from .controller import EcovacsController
+from . import EcovacsConfigEntry
+from .const import SUPPORTED_LIFESPANS
 from .entity import (
  CapabilityDevice,
  EcovacsCapabilityEntityDescription,
@@ -171,11 +170,11 @@ class EcovacsLifespanSensorEntityDescription(SensorEntityDescription):
 
 async def async_setup_entry(
  hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: EcovacsConfigEntry,
  async_add_entities: AddEntitiesCallback,
 ) -> None:
  """Add entities for passed config_entry in HA."""
- controller: EcovacsController = hass.data[DOMAIN][config_entry.entry_id]
+ controller = config_entry.runtime_data
 
  entities: list[EcovacsEntity] = get_supported_entitites(
  controller, EcovacsSensor, ENTITY_DESCRIPTIONS

homeassistant/components/ecovacs/switch.py

@@ -11,13 +11,11 @@
 from deebot_client.events import EnableEvent
 
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .controller import EcovacsController
+from . import EcovacsConfigEntry
 from .entity import (
  CapabilityDevice,
  EcovacsCapabilityEntityDescription,
@@ -121,11 +119,11 @@ class EcovacsSwitchEntityDescription(
 
 async def async_setup_entry(
  hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: EcovacsConfigEntry,
  async_add_entities: AddEntitiesCallback,
 ) -> None:
  """Add entities for passed config_entry in HA."""
- controller: EcovacsController = hass.data[DOMAIN][config_entry.entry_id]
+ controller = config_entry.runtime_data
  entities: list[EcovacsEntity] = get_supported_entitites(
  controller, EcovacsSwitchEntity, ENTITY_DESCRIPTIONS
  )

homeassistant/components/ecovacs/vacuum.py

@@ -23,15 +23,14 @@
  StateVacuumEntityDescription,
  VacuumEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ServiceValidationError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.icon import icon_for_battery_level
 from homeassistant.util import slugify
 
+from . import EcovacsConfigEntry
 from .const import DOMAIN
-from .controller import EcovacsController
 from .entity import EcovacsEntity
 from .util import get_name_key
 
@@ -43,11 +42,11 @@
 
 async def async_setup_entry(
  hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: EcovacsConfigEntry,
  async_add_entities: AddEntitiesCallback,
 ) -> None:
  """Set up the Ecovacs vacuums."""
- controller: EcovacsController = hass.data[DOMAIN][config_entry.entry_id]
+ controller = config_entry.runtime_data
  vacuums: list[EcovacsVacuum | EcovacsLegacyVacuum] = [
  EcovacsVacuum(device) for device in controller.devices(VacuumCapabilities)
  ]