Migrate aosmith to use runtime_data (#125538)

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

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 . import AOSmithConfigEntry
-from .const import DOMAIN
 
 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)

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 . import AOSmithConfigEntry
-from .const import DOMAIN
 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)

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 . import AOSmithConfigEntry
-from .const import DOMAIN
 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)