diff --git a/homeassistant/components/ecoforest/__init__.py b/homeassistant/components/ecoforest/__init__.py index 4d5aaa40576..e5350beba8e 100644 --- a/homeassistant/components/ecoforest/__init__.py +++ b/homeassistant/components/ecoforest/__init__.py @@ -11,20 +11,18 @@ from pyecoforest.exceptions import ( EcoforestConnectionError, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME, Platform from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryNotReady -from .const import DOMAIN -from .coordinator import EcoforestCoordinator +from .coordinator import EcoforestConfigEntry, EcoforestCoordinator PLATFORMS: list[Platform] = [Platform.NUMBER, Platform.SENSOR, Platform.SWITCH] _LOGGER = logging.getLogger(__name__) -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: EcoforestConfigEntry) -> bool: """Set up Ecoforest from a config entry.""" host = entry.data[CONF_HOST] @@ -41,20 +39,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: _LOGGER.error("Error communicating with device %s", host) raise ConfigEntryNotReady from err - coordinator = EcoforestCoordinator(hass, api) + coordinator = EcoforestCoordinator(hass, entry, api) await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + entry.runtime_data = coordinator 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: EcoforestConfigEntry) -> bool: """Unload a config entry.""" - if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - hass.data[DOMAIN].pop(entry.entry_id) - - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/ecoforest/coordinator.py b/homeassistant/components/ecoforest/coordinator.py index 3b04325bd50..603fde38388 100644 --- a/homeassistant/components/ecoforest/coordinator.py +++ b/homeassistant/components/ecoforest/coordinator.py @@ -6,6 +6,7 @@ from pyecoforest.api import EcoforestApi from pyecoforest.exceptions import EcoforestError from pyecoforest.models.device import Device +from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed @@ -13,16 +14,21 @@ from .const import POLLING_INTERVAL _LOGGER = logging.getLogger(__name__) +type EcoforestConfigEntry = ConfigEntry[EcoforestCoordinator] + class EcoforestCoordinator(DataUpdateCoordinator[Device]): """DataUpdateCoordinator to gather data from ecoforest device.""" - def __init__(self, hass: HomeAssistant, api: EcoforestApi) -> None: + def __init__( + self, hass: HomeAssistant, entry: EcoforestConfigEntry, api: EcoforestApi + ) -> None: """Initialize DataUpdateCoordinator.""" super().__init__( hass, _LOGGER, + config_entry=entry, name="ecoforest", update_interval=POLLING_INTERVAL, ) diff --git a/homeassistant/components/ecoforest/number.py b/homeassistant/components/ecoforest/number.py index db3275c1fcc..878c150343e 100644 --- a/homeassistant/components/ecoforest/number.py +++ b/homeassistant/components/ecoforest/number.py @@ -8,12 +8,10 @@ from dataclasses import dataclass from pyecoforest.models.device import Device from homeassistant.components.number import NumberEntity, NumberEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN -from .coordinator import EcoforestCoordinator +from .coordinator import EcoforestConfigEntry from .entity import EcoforestEntity @@ -38,11 +36,11 @@ NUMBER_ENTITIES = ( async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: EcoforestConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Ecoforest number platform.""" - coordinator: EcoforestCoordinator = hass.data[DOMAIN][config_entry.entry_id] + coordinator = config_entry.runtime_data entities = [ EcoforestNumberEntity(coordinator, description) diff --git a/homeassistant/components/ecoforest/sensor.py b/homeassistant/components/ecoforest/sensor.py index 997b02436cc..0babb476ab6 100644 --- a/homeassistant/components/ecoforest/sensor.py +++ b/homeassistant/components/ecoforest/sensor.py @@ -13,7 +13,6 @@ from homeassistant.components.sensor import ( SensorEntity, SensorEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( PERCENTAGE, UnitOfPressure, @@ -24,8 +23,7 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import StateType -from .const import DOMAIN -from .coordinator import EcoforestCoordinator +from .coordinator import EcoforestConfigEntry from .entity import EcoforestEntity _LOGGER = logging.getLogger(__name__) @@ -143,10 +141,12 @@ SENSOR_TYPES: tuple[EcoforestSensorEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: EcoforestConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up the Ecoforest sensor platform.""" - coordinator: EcoforestCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data entities = [ EcoforestSensor(coordinator, description) for description in SENSOR_TYPES diff --git a/homeassistant/components/ecoforest/switch.py b/homeassistant/components/ecoforest/switch.py index d643217bebc..de52248e751 100644 --- a/homeassistant/components/ecoforest/switch.py +++ b/homeassistant/components/ecoforest/switch.py @@ -10,12 +10,10 @@ from pyecoforest.api import EcoforestApi from pyecoforest.models.device import Device from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN -from .coordinator import EcoforestCoordinator +from .coordinator import EcoforestConfigEntry from .entity import EcoforestEntity @@ -39,11 +37,11 @@ SWITCH_TYPES: tuple[EcoforestSwitchEntityDescription, ...] = ( async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: EcoforestConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Ecoforest switch platform.""" - coordinator: EcoforestCoordinator = hass.data[DOMAIN][config_entry.entry_id] + coordinator = config_entry.runtime_data entities = [ EcoforestSwitchEntity(coordinator, description) for description in SWITCH_TYPES diff --git a/tests/components/ecoforest/conftest.py b/tests/components/ecoforest/conftest.py index 85bfff08bdf..8678cfd4d05 100644 --- a/tests/components/ecoforest/conftest.py +++ b/tests/components/ecoforest/conftest.py @@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, Mock, patch from pyecoforest.models.device import Alarm, Device, OperationMode, State import pytest -from homeassistant.components.ecoforest import DOMAIN +from homeassistant.components.ecoforest.const import DOMAIN from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME from homeassistant.core import HomeAssistant