From 7209b3c7d323e8681e2ec6734c7fc7d967bfd281 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Sun, 8 Sep 2024 22:05:48 +0200 Subject: [PATCH] Migrate aosmith to use runtime_data (#125538) --- homeassistant/components/aosmith/__init__.py | 13 ++++++------- homeassistant/components/aosmith/diagnostics.py | 8 +++----- homeassistant/components/aosmith/sensor.py | 10 +++++----- homeassistant/components/aosmith/water_heater.py | 10 +++++----- 4 files changed, 19 insertions(+), 22 deletions(-) diff --git a/homeassistant/components/aosmith/__init__.py b/homeassistant/components/aosmith/__init__.py index c42096cd3a7..dd60f69c4b9 100644 --- a/homeassistant/components/aosmith/__init__.py +++ b/homeassistant/components/aosmith/__init__.py @@ -16,6 +16,8 @@ from .coordinator import AOSmithEnergyCoordinator, AOSmithStatusCoordinator PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.WATER_HEATER] +type AOSmithConfigEntry = ConfigEntry[AOSmithData] + @dataclass class AOSmithData: @@ -26,7 +28,7 @@ class AOSmithData: energy_coordinator: AOSmithEnergyCoordinator -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: AOSmithConfigEntry) -> bool: """Set up A. O. Smith from a config entry.""" email = entry.data[CONF_EMAIL] password = entry.data[CONF_PASSWORD] @@ -55,7 +57,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: ) await energy_coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = AOSmithData( + entry.runtime_data = AOSmithData( client, status_coordinator, energy_coordinator, @@ -66,9 +68,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: AOSmithConfigEntry) -> 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/aosmith/diagnostics.py b/homeassistant/components/aosmith/diagnostics.py index 96b049b904f..94726731f75 100644 --- a/homeassistant/components/aosmith/diagnostics.py +++ b/homeassistant/components/aosmith/diagnostics.py @@ -5,11 +5,9 @@ from __future__ import annotations from typing import Any from homeassistant.components.diagnostics import async_redact_data -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from . import AOSmithData -from .const import DOMAIN +from . import AOSmithConfigEntry TO_REDACT = { "address", @@ -31,10 +29,10 @@ TO_REDACT = { async def async_get_config_entry_diagnostics( - hass: HomeAssistant, config_entry: ConfigEntry + hass: HomeAssistant, config_entry: AOSmithConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - data: AOSmithData = hass.data[DOMAIN][config_entry.entry_id] + data = config_entry.runtime_data all_device_info = await data.client.get_all_device_info() return async_redact_data(all_device_info, TO_REDACT) diff --git a/homeassistant/components/aosmith/sensor.py b/homeassistant/components/aosmith/sensor.py index e33c388af8b..89b383744e5 100644 --- a/homeassistant/components/aosmith/sensor.py +++ b/homeassistant/components/aosmith/sensor.py @@ -11,13 +11,11 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import UnitOfEnergy from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import AOSmithData -from .const import DOMAIN +from . import AOSmithConfigEntry from .coordinator import AOSmithEnergyCoordinator, AOSmithStatusCoordinator from .entity import AOSmithEnergyEntity, AOSmithStatusEntity @@ -49,10 +47,12 @@ HOT_WATER_STATUS_MAP: dict[HotWaterStatus, str] = { async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AOSmithConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up A. O. Smith sensor platform.""" - data: AOSmithData = hass.data[DOMAIN][entry.entry_id] + data = entry.runtime_data async_add_entities( AOSmithStatusSensorEntity(data.status_coordinator, description, junction_id) diff --git a/homeassistant/components/aosmith/water_heater.py b/homeassistant/components/aosmith/water_heater.py index dceba13ba34..f3dc8b3413f 100644 --- a/homeassistant/components/aosmith/water_heater.py +++ b/homeassistant/components/aosmith/water_heater.py @@ -12,14 +12,12 @@ from homeassistant.components.water_heater import ( WaterHeaterEntity, WaterHeaterEntityFeature, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import UnitOfTemperature from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import AOSmithData -from .const import DOMAIN +from . import AOSmithConfigEntry from .coordinator import AOSmithStatusCoordinator from .entity import AOSmithStatusEntity @@ -46,10 +44,12 @@ DEFAULT_OPERATION_MODE_PRIORITY = [ async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AOSmithConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up A. O. Smith water heater platform.""" - data: AOSmithData = hass.data[DOMAIN][entry.entry_id] + data = entry.runtime_data async_add_entities( AOSmithWaterHeaterEntity(data.status_coordinator, junction_id)