Convert powerwall to use entry.runtime_data (#121643)

homeassistant/components/powerwall/__init__.py

@@ -27,7 +27,12 @@ from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, Upda
 from homeassistant.util.network import is_ip_address
 
 from .const import DOMAIN, POWERWALL_API_CHANGED, POWERWALL_COORDINATOR, UPDATE_INTERVAL
-from .models import PowerwallBaseInfo, PowerwallData, PowerwallRuntimeData
+from .models import (
+    PowerwallBaseInfo,
+    PowerwallConfigEntry,
+    PowerwallData,
+    PowerwallRuntimeData,
+)
 
 CONFIG_SCHEMA = cv.removed(DOMAIN, raise_if_present=False)
 
@@ -115,7 +120,7 @@ class PowerwallDataManager:
         raise RuntimeError("unreachable")
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: PowerwallConfigEntry) -> bool:
     """Set up Tesla Powerwall from a config entry."""
     ip_address: str = entry.data[CONF_IP_ADDRESS]
 
@@ -176,7 +181,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     runtime_data[POWERWALL_COORDINATOR] = coordinator
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = runtime_data
+    entry.runtime_data = runtime_data
 
     await async_migrate_entity_unique_ids(hass, entry, base_info)
 
@@ -186,7 +191,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 
 async def async_migrate_entity_unique_ids(
-    hass: HomeAssistant, entry: ConfigEntry, base_info: PowerwallBaseInfo
+    hass: HomeAssistant, entry: PowerwallConfigEntry, base_info: PowerwallBaseInfo
 ) -> None:
     """Migrate old entity unique ids to use gateway_din."""
     old_base_unique_id = "_".join(base_info.serial_numbers)
@@ -276,21 +281,18 @@ async def _fetch_powerwall_data(power_wall: Powerwall) -> PowerwallData:
 
 
 @callback
-def async_last_update_was_successful(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+def async_last_update_was_successful(
+    hass: HomeAssistant, entry: PowerwallConfigEntry
+) -> bool:
     """Return True if the last update was successful."""
     return bool(
-        (domain_data := hass.data.get(DOMAIN))
-        and (entry_data := domain_data.get(entry.entry_id))
-        and (coordinator := entry_data.get(POWERWALL_COORDINATOR))
+        hasattr(entry, "runtime_data")
+        and (runtime_data := entry.runtime_data)
+        and (coordinator := runtime_data.get(POWERWALL_COORDINATOR))
         and coordinator.last_update_success
     )
 
 
 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)

homeassistant/components/powerwall/binary_sensor.py

@@ -8,13 +8,11 @@ from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
 from .entity import PowerWallEntity
-from .models import PowerwallRuntimeData
+from .models import PowerwallConfigEntry
 
 CONNECTED_GRID_STATUSES = {
     GridStatus.TRANSITION_TO_GRID,
@@ -24,11 +22,11 @@ CONNECTED_GRID_STATUSES = {
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: PowerwallConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the powerwall sensors."""
-    powerwall_data: PowerwallRuntimeData = hass.data[DOMAIN][config_entry.entry_id]
+    powerwall_data = entry.runtime_data
     async_add_entities(
         [
             sensor_class(powerwall_data)

homeassistant/components/powerwall/models.py

@@ -16,8 +16,11 @@ from tesla_powerwall import (
     SiteMasterResponse,
 )
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
+type PowerwallConfigEntry = ConfigEntry[PowerwallRuntimeData]
+
 
 @dataclass
 class PowerwallBaseInfo:

homeassistant/components/powerwall/sensor.py

@@ -14,7 +14,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     PERCENTAGE,
     EntityCategory,
@@ -28,9 +27,9 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN, POWERWALL_COORDINATOR
+from .const import POWERWALL_COORDINATOR
 from .entity import BatteryEntity, PowerWallEntity
-from .models import BatteryResponse, PowerwallRuntimeData
+from .models import BatteryResponse, PowerwallConfigEntry, PowerwallRuntimeData
 
 _METER_DIRECTION_EXPORT = "export"
 _METER_DIRECTION_IMPORT = "import"
@@ -219,11 +218,11 @@ BATTERY_INSTANT_SENSORS: list[PowerwallSensorEntityDescription] = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: PowerwallConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the powerwall sensors."""
-    powerwall_data: PowerwallRuntimeData = hass.data[DOMAIN][config_entry.entry_id]
+    powerwall_data = entry.runtime_data
     coordinator = powerwall_data[POWERWALL_COORDINATOR]
     assert coordinator is not None
     data = coordinator.data

homeassistant/components/powerwall/switch.py

@@ -5,15 +5,13 @@ from typing import Any
 from tesla_powerwall import GridStatus, IslandMode, PowerwallError
 
 from homeassistant.components.switch import SwitchDeviceClass, SwitchEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
 from .entity import PowerWallEntity
-from .models import PowerwallRuntimeData
+from .models import PowerwallConfigEntry, PowerwallRuntimeData
 
 OFF_GRID_STATUSES = {
     GridStatus.TRANSITION_TO_ISLAND,
@@ -23,12 +21,11 @@ OFF_GRID_STATUSES = {
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: PowerwallConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Powerwall switch platform from Powerwall resources."""
-    powerwall_data: PowerwallRuntimeData = hass.data[DOMAIN][config_entry.entry_id]
-    async_add_entities([PowerwallOffGridEnabledEntity(powerwall_data)])
+    async_add_entities([PowerwallOffGridEnabledEntity(entry.runtime_data)])
 
 
 class PowerwallOffGridEnabledEntity(PowerWallEntity, SwitchEntity):