diff --git a/homeassistant/components/discovergy/__init__.py b/homeassistant/components/discovergy/__init__.py index 0d38182da5d..974441f3899 100644 --- a/homeassistant/components/discovergy/__init__.py +++ b/homeassistant/components/discovergy/__init__.py @@ -12,16 +12,15 @@ from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from homeassistant.helpers.httpx_client import get_async_client -from .const import DOMAIN from .coordinator import DiscovergyUpdateCoordinator PLATFORMS = [Platform.SENSOR] +DiscovergyConfigEntry = ConfigEntry[list[DiscovergyUpdateCoordinator]] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + +async def async_setup_entry(hass: HomeAssistant, entry: DiscovergyConfigEntry) -> bool: """Set up Discovergy from a config entry.""" - hass.data.setdefault(DOMAIN, {}) - client = Discovergy( email=entry.data[CONF_EMAIL], password=entry.data[CONF_PASSWORD], @@ -53,7 +52,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: await coordinator.async_config_entry_first_refresh() coordinators.append(coordinator) - hass.data[DOMAIN][entry.entry_id] = coordinators + entry.runtime_data = coordinators await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) entry.async_on_unload(entry.add_update_listener(async_reload_entry)) @@ -63,11 +62,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload a config entry.""" - unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) - - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) async def async_reload_entry(hass: HomeAssistant, entry: ConfigEntry) -> None: diff --git a/homeassistant/components/discovergy/diagnostics.py b/homeassistant/components/discovergy/diagnostics.py index 15676da9888..3857404db81 100644 --- a/homeassistant/components/discovergy/diagnostics.py +++ b/homeassistant/components/discovergy/diagnostics.py @@ -6,11 +6,9 @@ from dataclasses import asdict from typing import Any from homeassistant.components.diagnostics import async_redact_data -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from .const import DOMAIN -from .coordinator import DiscovergyUpdateCoordinator +from . import DiscovergyConfigEntry TO_REDACT_METER = { "serial_number", @@ -22,14 +20,13 @@ TO_REDACT_METER = { async def async_get_config_entry_diagnostics( - hass: HomeAssistant, entry: ConfigEntry + hass: HomeAssistant, entry: DiscovergyConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" flattened_meter: list[dict] = [] last_readings: dict[str, dict] = {} - coordinators: list[DiscovergyUpdateCoordinator] = hass.data[DOMAIN][entry.entry_id] - for coordinator in coordinators: + for coordinator in entry.runtime_data: # make a dict of meter data and redact some data flattened_meter.append( async_redact_data(asdict(coordinator.meter), TO_REDACT_METER) diff --git a/homeassistant/components/discovergy/sensor.py b/homeassistant/components/discovergy/sensor.py index 0a820917821..531904c8740 100644 --- a/homeassistant/components/discovergy/sensor.py +++ b/homeassistant/components/discovergy/sensor.py @@ -12,7 +12,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( EntityCategory, UnitOfElectricPotential, @@ -25,6 +24,7 @@ from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity +from . import DiscovergyConfigEntry from .const import DOMAIN, MANUFACTURER from .coordinator import DiscovergyUpdateCoordinator @@ -163,13 +163,13 @@ ADDITIONAL_SENSORS: tuple[DiscovergySensorEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: DiscovergyConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up the Discovergy sensors.""" - coordinators: list[DiscovergyUpdateCoordinator] = hass.data[DOMAIN][entry.entry_id] - entities: list[DiscovergySensor] = [] - for coordinator in coordinators: + for coordinator in entry.runtime_data: sensors: tuple[DiscovergySensorEntityDescription, ...] = () # select sensor descriptions based on meter type and combine with additional sensors diff --git a/tests/components/discovergy/conftest.py b/tests/components/discovergy/conftest.py index d3ab3b831f0..913e33f6367 100644 --- a/tests/components/discovergy/conftest.py +++ b/tests/components/discovergy/conftest.py @@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, patch from pydiscovergy.models import Reading import pytest -from homeassistant.components.discovergy import DOMAIN +from homeassistant.components.discovergy.const import DOMAIN from homeassistant.const import CONF_EMAIL, CONF_PASSWORD from homeassistant.core import HomeAssistant from homeassistant.setup import async_setup_component