Use runtime_data in balboa (#129035)

homeassistant/components/balboa/__init__.py

@@ -14,7 +14,7 @@ from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.event import async_track_time_interval
 import homeassistant.util.dt as dt_util
 
-from .const import CONF_SYNC_TIME, DEFAULT_SYNC_TIME, DOMAIN
+from .const import CONF_SYNC_TIME, DEFAULT_SYNC_TIME
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -30,8 +30,10 @@ PLATFORMS = [
 KEEP_ALIVE_INTERVAL = timedelta(minutes=1)
 SYNC_TIME_INTERVAL = timedelta(hours=1)
 
+type BalboaConfigEntry = ConfigEntry[SpaClient]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: BalboaConfigEntry) -> bool:
     """Set up Balboa Spa from a config entry."""
     host = entry.data[CONF_HOST]
 
@@ -44,41 +46,34 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         _LOGGER.error("Failed to get spa info at %s", host)
         raise ConfigEntryNotReady("Unable to configure")
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = spa
+    entry.runtime_data = spa
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     await async_setup_time_sync(hass, entry)
     entry.async_on_unload(entry.add_update_listener(update_listener))
+    entry.async_on_unload(spa.disconnect)
 
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: BalboaConfigEntry) -> bool:
     """Unload a config entry."""
-    _LOGGER.debug("Disconnecting from spa")
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    spa: SpaClient = hass.data[DOMAIN][entry.entry_id]
-
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    await spa.disconnect()
-
-    return unload_ok
 
 
-async def update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def update_listener(hass: HomeAssistant, entry: BalboaConfigEntry) -> None:
     """Handle options update."""
     await hass.config_entries.async_reload(entry.entry_id)
 
 
-async def async_setup_time_sync(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def async_setup_time_sync(hass: HomeAssistant, entry: BalboaConfigEntry) -> None:
     """Set up the time sync."""
     if not entry.options.get(CONF_SYNC_TIME, DEFAULT_SYNC_TIME):
         return
 
     _LOGGER.debug("Setting up daily time sync")
-    spa: SpaClient = hass.data[DOMAIN][entry.entry_id]
+    spa = entry.runtime_data
 
     async def sync_time(now: datetime) -> None:
         now = dt_util.as_local(now)

homeassistant/components/balboa/binary_sensor.py

@@ -12,19 +12,20 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import BalboaConfigEntry
 from .entity import BalboaEntity
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: BalboaConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the spa's binary sensors."""
-    spa: SpaClient = hass.data[DOMAIN][entry.entry_id]
+    spa = entry.runtime_data
     entities = [
         BalboaBinarySensorEntity(spa, description)
         for description in BINARY_SENSOR_DESCRIPTIONS

homeassistant/components/balboa/climate.py

@@ -14,7 +14,6 @@ from homeassistant.components.climate import (
     HVACAction,
     HVACMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_TEMPERATURE,
     PRECISION_HALVES,
@@ -24,6 +23,7 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from . import BalboaConfigEntry
 from .const import DOMAIN
 from .entity import BalboaEntity
 
@@ -45,10 +45,12 @@ TEMPERATURE_UNIT_MAP = {
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: BalboaConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the spa climate entity."""
-    async_add_entities([BalboaClimateEntity(hass.data[DOMAIN][entry.entry_id])])
+    async_add_entities([BalboaClimateEntity(entry.runtime_data)])
 
 
 class BalboaClimateEntity(BalboaEntity, ClimateEntity):

homeassistant/components/balboa/fan.py

@@ -5,11 +5,10 @@ from __future__ import annotations
 import math
 from typing import Any, cast
 
-from pybalboa import SpaClient, SpaControl
+from pybalboa import SpaControl
 from pybalboa.enums import OffOnState, UnknownState
 
 from homeassistant.components.fan import FanEntity, FanEntityFeature
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.percentage import (
@@ -17,15 +16,17 @@ from homeassistant.util.percentage import (
     ranged_value_to_percentage,
 )
 
-from .const import DOMAIN
+from . import BalboaConfigEntry
 from .entity import BalboaEntity
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: BalboaConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the spa's pumps."""
-    spa: SpaClient = hass.data[DOMAIN][entry.entry_id]
+    spa = entry.runtime_data
     async_add_entities(BalboaPumpFanEntity(control) for control in spa.pumps)
 
 

homeassistant/components/balboa/light.py

@@ -4,23 +4,24 @@ from __future__ import annotations
 
 from typing import Any, cast
 
-from pybalboa import SpaClient, SpaControl
+from pybalboa import SpaControl
 from pybalboa.enums import OffOnState, UnknownState
 
 from homeassistant.components.light import ColorMode, LightEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import BalboaConfigEntry
 from .entity import BalboaEntity
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: BalboaConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the spa's lights."""
-    spa: SpaClient = hass.data[DOMAIN][entry.entry_id]
+    spa = entry.runtime_data
     async_add_entities(BalboaLightEntity(control) for control in spa.lights)
 
 

homeassistant/components/balboa/select.py

@@ -1,22 +1,23 @@
 """Support for Spa Client selects."""
 
-from pybalboa import SpaClient, SpaControl
+from pybalboa import SpaControl
 from pybalboa.enums import LowHighRange
 
 from homeassistant.components.select import SelectEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import BalboaConfigEntry
 from .entity import BalboaEntity
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: BalboaConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the spa select entity."""
-    spa: SpaClient = hass.data[DOMAIN][entry.entry_id]
+    spa = entry.runtime_data
     async_add_entities([BalboaTempRangeSelectEntity(spa.temperature_range)])
 
 

tests/components/balboa/__init__.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 from unittest.mock import MagicMock
 
-from homeassistant.components.balboa import CONF_SYNC_TIME, DOMAIN
+from homeassistant.components.balboa.const import CONF_SYNC_TIME, DOMAIN
 from homeassistant.const import CONF_HOST
 from homeassistant.core import HomeAssistant, State