Use runtime_data in ecoforest (#136689)

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 EcoforestConfigEntry, EcoforestCoordinator
-from .coordinator import 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):
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok

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,
         )

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 EcoforestConfigEntry
-from .coordinator import EcoforestCoordinator
 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)

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 EcoforestConfigEntry
-from .coordinator import EcoforestCoordinator
 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

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 EcoforestConfigEntry
-from .coordinator import EcoforestCoordinator
 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

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