Store runtime data in entry in Ecovacs (#116445)

homeassistant/components/ecovacs/__init__.py

@@ -37,6 +37,7 @@ PLATFORMS = [
     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):
+    return 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

homeassistant/components/ecovacs/binary_sensor.py

@@ -11,13 +11,11 @@ from homeassistant.components.binary_sensor import (
     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 . import EcovacsConfigEntry
-from .controller import EcovacsController
 from .entity import (
     CapabilityDevice,
     EcovacsCapabilityEntityDescription,
@@ -52,13 +50,14 @@ ENTITY_DESCRIPTIONS: tuple[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.capabilities import (
 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 . import EcovacsConfigEntry
-from .controller import EcovacsController
+from .const import SUPPORTED_LIFESPANS
 from .entity import (
     CapabilityDevice,
     EcovacsCapabilityEntityDescription,
@@ -66,11 +65,11 @@ LIFESPAN_ENTITY_DESCRIPTIONS = tuple(
 
 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 @@ class EcovacsController:
         """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 @@ class EcovacsController:
                         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 @@ class EcovacsController:
         """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 @@ class EcovacsController:
         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 typing import Any
 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 . import EcovacsConfigEntry
-from .controller import EcovacsController
+from .const import CONF_OVERRIDE_MQTT_URL, CONF_OVERRIDE_REST_URL
 
 REDACT_CONFIG = {
     CONF_USERNAME,
@@ -25,10 +24,10 @@ REDACT_DEVICE = {"did", CONF_NAME, "homeId"}
 
 
 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.device import Device
 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 . import EcovacsConfigEntry
-from .controller import EcovacsController
 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.device import Device
 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 . import EcovacsConfigEntry
-from .controller import EcovacsController
 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 @@ from homeassistant.components.lawn_mower import (
     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 . import EcovacsConfigEntry
-from .controller import EcovacsController
 from .entity import EcovacsEntity
 
 _LOGGER = logging.getLogger(__name__)
@@ -38,11 +36,11 @@ _STATE_TO_MOWER_STATE = {
 
 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.capabilities import Capabilities, CapabilitySet, VacuumCapabi
 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 . import EcovacsConfigEntry
-from .controller import EcovacsController
 from .entity import (
     CapabilityDevice,
     EcovacsCapabilityEntityDescription,
@@ -70,11 +68,11 @@ ENTITY_DESCRIPTIONS: tuple[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.device import Device
 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 . import EcovacsConfigEntry
-from .controller import EcovacsController
 from .entity import (
     CapabilityDevice,
     EcovacsCapabilityEntityDescription,
@@ -62,11 +60,11 @@ ENTITY_DESCRIPTIONS: tuple[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 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     AREA_SQUARE_METERS,
     ATTR_BATTERY_LEVEL,
@@ -37,8 +36,8 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
-from .const import DOMAIN, SUPPORTED_LIFESPANS
+from . import EcovacsConfigEntry
-from .controller import EcovacsController
+from .const import SUPPORTED_LIFESPANS
 from .entity import (
     CapabilityDevice,
     EcovacsCapabilityEntityDescription,
@@ -171,11 +170,11 @@ LIFESPAN_ENTITY_DESCRIPTIONS = tuple(
 
 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.capabilities import (
 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 . import EcovacsConfigEntry
-from .controller import EcovacsController
 from .entity import (
     CapabilityDevice,
     EcovacsCapabilityEntityDescription,
@@ -121,11 +119,11 @@ ENTITY_DESCRIPTIONS: tuple[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 @@ from homeassistant.components.vacuum import (
     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 @@ ATTR_COMPONENT_PREFIX = "component_"
 
 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)
     ]

tests/components/ecovacs/conftest.py

@@ -156,8 +156,6 @@ async def init_integration(
 
 
 @pytest.fixture
-def controller(
+def controller(init_integration: MockConfigEntry) -> EcovacsController:
-    hass: HomeAssistant, init_integration: MockConfigEntry
-) -> EcovacsController:
     """Get the controller for the config entry."""
-    return hass.data[DOMAIN][init_integration.entry_id]
+    return init_integration.runtime_data

tests/components/ecovacs/test_init.py

@@ -20,21 +20,34 @@ from .const import IMPORT_DATA
 from tests.common import MockConfigEntry
 
 
-@pytest.mark.usefixtures("init_integration")
+@pytest.mark.usefixtures(
+    "mock_authenticator", "mock_mqtt_client", "mock_device_execute"
+)
 async def test_load_unload_config_entry(
     hass: HomeAssistant,
-    init_integration: MockConfigEntry,
+    mock_config_entry: MockConfigEntry,
 ) -> None:
     """Test loading and unloading the integration."""
-    mock_config_entry = init_integration
+    with patch(
+        "homeassistant.components.ecovacs.EcovacsController",
+        autospec=True,
+    ):
+        mock_config_entry.add_to_hass(hass)
+
+        await hass.config_entries.async_setup(mock_config_entry.entry_id)
+        await hass.async_block_till_done()
+
         assert mock_config_entry.state is ConfigEntryState.LOADED
-    assert DOMAIN in hass.data
+        assert DOMAIN not in hass.data
+        controller = mock_config_entry.runtime_data
+        assert isinstance(controller, EcovacsController)
+        controller.initialize.assert_called_once()
 
         await hass.config_entries.async_unload(mock_config_entry.entry_id)
         await hass.async_block_till_done()
+        controller.teardown.assert_called_once()
 
         assert mock_config_entry.state is ConfigEntryState.NOT_LOADED
-    assert DOMAIN not in hass.data
 
 
 @pytest.fixture