diff --git a/homeassistant/components/suez_water/__init__.py b/homeassistant/components/suez_water/__init__.py index 06f503b85c26d8..cbaac91264248a 100644 --- a/homeassistant/components/suez_water/__init__.py +++ b/homeassistant/components/suez_water/__init__.py @@ -2,32 +2,27 @@ from __future__ import annotations -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from .const import DOMAIN -from .coordinator import SuezWaterCoordinator +from .coordinator import SuezWaterConfigEntry, SuezWaterCoordinator PLATFORMS: list[Platform] = [Platform.SENSOR] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: SuezWaterConfigEntry) -> bool: """Set up Suez Water from a config entry.""" coordinator = SuezWaterCoordinator(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: SuezWaterConfigEntry) -> 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/suez_water/coordinator.py b/homeassistant/components/suez_water/coordinator.py index 224929c606e595..72da68c0f5d19e 100644 --- a/homeassistant/components/suez_water/coordinator.py +++ b/homeassistant/components/suez_water/coordinator.py @@ -37,13 +37,16 @@ class SuezWaterData: price: float +type SuezWaterConfigEntry = ConfigEntry[SuezWaterCoordinator] + + class SuezWaterCoordinator(DataUpdateCoordinator[SuezWaterData]): """Suez water coordinator.""" _suez_client: SuezClient - config_entry: ConfigEntry + config_entry: SuezWaterConfigEntry - def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry) -> None: + def __init__(self, hass: HomeAssistant, config_entry: SuezWaterConfigEntry) -> None: """Initialize suez water coordinator.""" super().__init__( hass, diff --git a/homeassistant/components/suez_water/quality_scale.yaml b/homeassistant/components/suez_water/quality_scale.yaml index 474340a14893cc..399c0b73a5af81 100644 --- a/homeassistant/components/suez_water/quality_scale.yaml +++ b/homeassistant/components/suez_water/quality_scale.yaml @@ -4,9 +4,7 @@ rules: test-before-configure: done unique-config-entry: done config-flow-test-coverage: done - runtime-data: - status: todo - comment: coordinator is created during setup, should be stored in runtime_data + runtime-data: done test-before-setup: done appropriate-polling: done entity-unique-id: done diff --git a/homeassistant/components/suez_water/sensor.py b/homeassistant/components/suez_water/sensor.py index 2ba699a9af14d3..e4e53dd7f6db80 100644 --- a/homeassistant/components/suez_water/sensor.py +++ b/homeassistant/components/suez_water/sensor.py @@ -13,7 +13,6 @@ SensorEntity, SensorEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CURRENCY_EURO, UnitOfVolume from homeassistant.core import HomeAssistant from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo @@ -21,7 +20,7 @@ from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import CONF_COUNTER_ID, DOMAIN -from .coordinator import SuezWaterCoordinator, SuezWaterData +from .coordinator import SuezWaterConfigEntry, SuezWaterCoordinator, SuezWaterData @dataclass(frozen=True, kw_only=True) @@ -53,11 +52,11 @@ class SuezWaterSensorEntityDescription(SensorEntityDescription): async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: SuezWaterConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Suez Water sensor from a config entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data counter_id = entry.data[CONF_COUNTER_ID] async_add_entities(