diff --git a/homeassistant/components/ecovacs/__init__.py b/homeassistant/components/ecovacs/__init__.py index ca4579a31b2..e4924b57641 100644 --- a/homeassistant/components/ecovacs/__init__.py +++ b/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): - 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) diff --git a/homeassistant/components/ecovacs/binary_sensor.py b/homeassistant/components/ecovacs/binary_sensor.py index cc401cc3ca0..f6e3e34aaa4 100644 --- a/homeassistant/components/ecovacs/binary_sensor.py +++ b/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 .controller import EcovacsController +from . import EcovacsConfigEntry 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 + ) ) diff --git a/homeassistant/components/ecovacs/button.py b/homeassistant/components/ecovacs/button.py index 27f729a1ae0..14fd54df5a0 100644 --- a/homeassistant/components/ecovacs/button.py +++ b/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 .controller import EcovacsController +from . import EcovacsConfigEntry +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 ) diff --git a/homeassistant/components/ecovacs/controller.py b/homeassistant/components/ecovacs/controller.py index 6b6fe3128dd..690f4e56cc9 100644 --- a/homeassistant/components/ecovacs/controller.py +++ b/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 diff --git a/homeassistant/components/ecovacs/diagnostics.py b/homeassistant/components/ecovacs/diagnostics.py index 9340841223e..50b59b90860 100644 --- a/homeassistant/components/ecovacs/diagnostics.py +++ b/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 .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 @@ 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) } diff --git a/homeassistant/components/ecovacs/event.py b/homeassistant/components/ecovacs/event.py index fb4c25c7559..9e4dde00b54 100644 --- a/homeassistant/components/ecovacs/event.py +++ b/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 .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) ) diff --git a/homeassistant/components/ecovacs/image.py b/homeassistant/components/ecovacs/image.py index 82e20e19732..1e94dc856ee 100644 --- a/homeassistant/components/ecovacs/image.py +++ b/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 .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 diff --git a/homeassistant/components/ecovacs/lawn_mower.py b/homeassistant/components/ecovacs/lawn_mower.py index 1b13d50cc0c..2561fe22217 100644 --- a/homeassistant/components/ecovacs/lawn_mower.py +++ b/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 .controller import EcovacsController +from . import EcovacsConfigEntry 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) ] diff --git a/homeassistant/components/ecovacs/number.py b/homeassistant/components/ecovacs/number.py index e53f7e6aae0..bd8ce50aadb 100644 --- a/homeassistant/components/ecovacs/number.py +++ b/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 .controller import EcovacsController +from . import EcovacsConfigEntry 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 ) diff --git a/homeassistant/components/ecovacs/select.py b/homeassistant/components/ecovacs/select.py index 01d4c5aae6b..4caa6327bb3 100644 --- a/homeassistant/components/ecovacs/select.py +++ b/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 .controller import EcovacsController +from . import EcovacsConfigEntry 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 ) diff --git a/homeassistant/components/ecovacs/sensor.py b/homeassistant/components/ecovacs/sensor.py index 92d1b10a614..e9229781827 100644 --- a/homeassistant/components/ecovacs/sensor.py +++ b/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 .controller import EcovacsController +from . import EcovacsConfigEntry +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 diff --git a/homeassistant/components/ecovacs/switch.py b/homeassistant/components/ecovacs/switch.py index 0d2f8f2024f..25ecb53e278 100644 --- a/homeassistant/components/ecovacs/switch.py +++ b/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 .controller import EcovacsController +from . import EcovacsConfigEntry 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 ) diff --git a/homeassistant/components/ecovacs/vacuum.py b/homeassistant/components/ecovacs/vacuum.py index 0e990645d7c..5c898694cbb 100644 --- a/homeassistant/components/ecovacs/vacuum.py +++ b/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) ] diff --git a/tests/components/ecovacs/conftest.py b/tests/components/ecovacs/conftest.py index 1a313957c3e..d4333f65dc4 100644 --- a/tests/components/ecovacs/conftest.py +++ b/tests/components/ecovacs/conftest.py @@ -156,8 +156,6 @@ async def init_integration( @pytest.fixture -def controller( - hass: HomeAssistant, init_integration: MockConfigEntry -) -> EcovacsController: +def controller(init_integration: MockConfigEntry) -> EcovacsController: """Get the controller for the config entry.""" - return hass.data[DOMAIN][init_integration.entry_id] + return init_integration.runtime_data diff --git a/tests/components/ecovacs/test_init.py b/tests/components/ecovacs/test_init.py index c27da2196b1..752276015d3 100644 --- a/tests/components/ecovacs/test_init.py +++ b/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 - assert mock_config_entry.state is ConfigEntryState.LOADED - assert DOMAIN in hass.data + with patch( + "homeassistant.components.ecovacs.EcovacsController", + autospec=True, + ): + mock_config_entry.add_to_hass(hass) - await hass.config_entries.async_unload(mock_config_entry.entry_id) - await hass.async_block_till_done() + await hass.config_entries.async_setup(mock_config_entry.entry_id) + await hass.async_block_till_done() - assert mock_config_entry.state is ConfigEntryState.NOT_LOADED - assert DOMAIN not in hass.data + assert mock_config_entry.state is ConfigEntryState.LOADED + 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 @pytest.fixture