diff --git a/homeassistant/components/iotawatt/__init__.py b/homeassistant/components/iotawatt/__init__.py index 8f35d4e0796..1dc38ba01c6 100644 --- a/homeassistant/components/iotawatt/__init__.py +++ b/homeassistant/components/iotawatt/__init__.py @@ -1,26 +1,22 @@ """The iotawatt integration.""" -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from .const import DOMAIN -from .coordinator import IotawattUpdater +from .coordinator import IotawattConfigEntry, IotawattUpdater PLATFORMS = [Platform.SENSOR] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: IotawattConfigEntry) -> bool: """Set up iotawatt from a config entry.""" coordinator = IotawattUpdater(hass, entry) 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: IotawattConfigEntry) -> 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/iotawatt/coordinator.py b/homeassistant/components/iotawatt/coordinator.py index 13802ebdd76..48d55dad818 100644 --- a/homeassistant/components/iotawatt/coordinator.py +++ b/homeassistant/components/iotawatt/coordinator.py @@ -21,14 +21,16 @@ _LOGGER = logging.getLogger(__name__) # Matches iotwatt data log interval REQUEST_REFRESH_DEFAULT_COOLDOWN = 5 +type IotawattConfigEntry = ConfigEntry[IotawattUpdater] + class IotawattUpdater(DataUpdateCoordinator): """Class to manage fetching update data from the IoTaWatt Energy Device.""" api: Iotawatt | None = None - config_entry: ConfigEntry + config_entry: IotawattConfigEntry - def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None: + def __init__(self, hass: HomeAssistant, entry: IotawattConfigEntry) -> None: """Initialize IotaWattUpdater object.""" super().__init__( hass=hass, diff --git a/homeassistant/components/iotawatt/sensor.py b/homeassistant/components/iotawatt/sensor.py index f5210f7fbba..591397ad6e7 100644 --- a/homeassistant/components/iotawatt/sensor.py +++ b/homeassistant/components/iotawatt/sensor.py @@ -14,7 +14,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( PERCENTAGE, UnitOfApparentPower, @@ -31,8 +30,8 @@ from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.util import dt as dt_util -from .const import DOMAIN, VOLT_AMPERE_REACTIVE, VOLT_AMPERE_REACTIVE_HOURS -from .coordinator import IotawattUpdater +from .const import VOLT_AMPERE_REACTIVE, VOLT_AMPERE_REACTIVE_HOURS +from .coordinator import IotawattConfigEntry, IotawattUpdater _LOGGER = logging.getLogger(__name__) @@ -113,11 +112,11 @@ ENTITY_DESCRIPTION_KEY_MAP: dict[str, IotaWattSensorEntityDescription] = { async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: IotawattConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Add sensors for passed config_entry in HA.""" - coordinator: IotawattUpdater = hass.data[DOMAIN][config_entry.entry_id] + coordinator = config_entry.runtime_data created = set() @callback diff --git a/tests/components/iotawatt/conftest.py b/tests/components/iotawatt/conftest.py index 9380154b53e..3b30783494e 100644 --- a/tests/components/iotawatt/conftest.py +++ b/tests/components/iotawatt/conftest.py @@ -5,7 +5,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from homeassistant.components.iotawatt import DOMAIN +from homeassistant.components.iotawatt.const import DOMAIN from homeassistant.core import HomeAssistant from tests.common import MockConfigEntry