diff --git a/homeassistant/components/balboa/__init__.py b/homeassistant/components/balboa/__init__.py index 7e220bd46f8..7838db16820 100644 --- a/homeassistant/components/balboa/__init__.py +++ b/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") - 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 + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) -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) diff --git a/homeassistant/components/balboa/binary_sensor.py b/homeassistant/components/balboa/binary_sensor.py index d3352208cd9..b8c62ce8abf 100644 --- a/homeassistant/components/balboa/binary_sensor.py +++ b/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 diff --git a/homeassistant/components/balboa/climate.py b/homeassistant/components/balboa/climate.py index 8cd9e93e539..d27fd459676 100644 --- a/homeassistant/components/balboa/climate.py +++ b/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): diff --git a/homeassistant/components/balboa/fan.py b/homeassistant/components/balboa/fan.py index bf7425f0e64..67c1d9a9a62 100644 --- a/homeassistant/components/balboa/fan.py +++ b/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) diff --git a/homeassistant/components/balboa/light.py b/homeassistant/components/balboa/light.py index 5dc8d48ef9d..21e4dfc5e08 100644 --- a/homeassistant/components/balboa/light.py +++ b/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) diff --git a/homeassistant/components/balboa/select.py b/homeassistant/components/balboa/select.py index 9c3074350c5..e88e40ab063 100644 --- a/homeassistant/components/balboa/select.py +++ b/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)]) diff --git a/tests/components/balboa/__init__.py b/tests/components/balboa/__init__.py index a27293e955f..2cb100e3642 100644 --- a/tests/components/balboa/__init__.py +++ b/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