From a12bc70543e7f904da356a766634ae840a3f21e3 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Tue, 20 May 2025 09:15:26 +0200 Subject: [PATCH] Use runtime_data in smarttub (#145279) --- homeassistant/components/smarttub/__init__.py | 19 ++++------- .../components/smarttub/binary_sensor.py | 32 ++++++++++++------ homeassistant/components/smarttub/climate.py | 13 +++++--- homeassistant/components/smarttub/const.py | 2 -- .../components/smarttub/controller.py | 33 ++++++++++--------- homeassistant/components/smarttub/entity.py | 19 ++++++++--- homeassistant/components/smarttub/light.py | 19 +++++------ homeassistant/components/smarttub/sensor.py | 21 +++++++----- homeassistant/components/smarttub/switch.py | 13 +++++--- tests/components/smarttub/test_init.py | 5 ++- 10 files changed, 99 insertions(+), 77 deletions(-) diff --git a/homeassistant/components/smarttub/__init__.py b/homeassistant/components/smarttub/__init__.py index 8406fdc4c2f..178fd9a70e2 100644 --- a/homeassistant/components/smarttub/__init__.py +++ b/homeassistant/components/smarttub/__init__.py @@ -1,11 +1,9 @@ """SmartTub integration.""" -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from .const import DOMAIN, SMARTTUB_CONTROLLER -from .controller import SmartTubController +from .controller import SmartTubConfigEntry, SmartTubController PLATFORMS = [ Platform.BINARY_SENSOR, @@ -16,26 +14,21 @@ PLATFORMS = [ ] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: SmartTubConfigEntry) -> bool: """Set up a smarttub config entry.""" controller = SmartTubController(hass) - hass.data.setdefault(DOMAIN, {}) - hass.data[DOMAIN][entry.entry_id] = { - SMARTTUB_CONTROLLER: controller, - } if not await controller.async_setup_entry(entry): return False + entry.runtime_data = controller + await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: SmartTubConfigEntry) -> bool: """Remove a smarttub 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) diff --git a/homeassistant/components/smarttub/binary_sensor.py b/homeassistant/components/smarttub/binary_sensor.py index 2e8792140b0..a120650e84b 100644 --- a/homeassistant/components/smarttub/binary_sensor.py +++ b/homeassistant/components/smarttub/binary_sensor.py @@ -2,20 +2,23 @@ from __future__ import annotations -from smarttub import SpaError, SpaReminder +from typing import Any + +from smarttub import Spa, SpaError, SpaReminder import voluptuous as vol from homeassistant.components.binary_sensor import ( BinarySensorDeviceClass, BinarySensorEntity, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers import entity_platform from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.typing import VolDictType +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from .const import ATTR_ERRORS, ATTR_REMINDERS, DOMAIN, SMARTTUB_CONTROLLER +from .const import ATTR_ERRORS, ATTR_REMINDERS +from .controller import SmartTubConfigEntry from .entity import SmartTubEntity, SmartTubSensorBase # whether the reminder has been snoozed (bool) @@ -44,12 +47,12 @@ SNOOZE_REMINDER_SCHEMA: VolDictType = { async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: SmartTubConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up binary sensor entities for the binary sensors in the tub.""" - controller = hass.data[DOMAIN][entry.entry_id][SMARTTUB_CONTROLLER] + controller = entry.runtime_data entities: list[BinarySensorEntity] = [] for spa in controller.spas: @@ -83,7 +86,9 @@ class SmartTubOnline(SmartTubSensorBase, BinarySensorEntity): # This seems to be very noisy and not generally useful, so disable by default. _attr_entity_registry_enabled_default = False - def __init__(self, coordinator, spa): + def __init__( + self, coordinator: DataUpdateCoordinator[dict[str, Any]], spa: Spa + ) -> None: """Initialize the entity.""" super().__init__(coordinator, spa, "Online", "online") @@ -98,7 +103,12 @@ class SmartTubReminder(SmartTubEntity, BinarySensorEntity): _attr_device_class = BinarySensorDeviceClass.PROBLEM - def __init__(self, coordinator, spa, reminder): + def __init__( + self, + coordinator: DataUpdateCoordinator[dict[str, Any]], + spa: Spa, + reminder: SpaReminder, + ) -> None: """Initialize the entity.""" super().__init__( coordinator, @@ -119,7 +129,7 @@ class SmartTubReminder(SmartTubEntity, BinarySensorEntity): return self.reminder.remaining_days == 0 @property - def extra_state_attributes(self): + def extra_state_attributes(self) -> dict[str, Any]: """Return the state attributes.""" return { ATTR_REMINDER_SNOOZED: self.reminder.snoozed, @@ -145,7 +155,9 @@ class SmartTubError(SmartTubEntity, BinarySensorEntity): _attr_device_class = BinarySensorDeviceClass.PROBLEM - def __init__(self, coordinator, spa): + def __init__( + self, coordinator: DataUpdateCoordinator[dict[str, Any]], spa: Spa + ) -> None: """Initialize the entity.""" super().__init__( coordinator, @@ -167,7 +179,7 @@ class SmartTubError(SmartTubEntity, BinarySensorEntity): return self.error is not None @property - def extra_state_attributes(self): + def extra_state_attributes(self) -> dict[str, Any]: """Return the state attributes.""" if (error := self.error) is None: return {} diff --git a/homeassistant/components/smarttub/climate.py b/homeassistant/components/smarttub/climate.py index f5759f32fa3..7e79ce0eb12 100644 --- a/homeassistant/components/smarttub/climate.py +++ b/homeassistant/components/smarttub/climate.py @@ -14,13 +14,14 @@ from homeassistant.components.climate import ( HVACAction, HVACMode, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from homeassistant.util.unit_conversion import TemperatureConverter -from .const import DEFAULT_MAX_TEMP, DEFAULT_MIN_TEMP, DOMAIN, SMARTTUB_CONTROLLER +from .const import DEFAULT_MAX_TEMP, DEFAULT_MIN_TEMP +from .controller import SmartTubConfigEntry from .entity import SmartTubEntity PRESET_DAY = "day" @@ -43,12 +44,12 @@ HVAC_ACTIONS = { async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: SmartTubConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up climate entity for the thermostat in the tub.""" - controller = hass.data[DOMAIN][entry.entry_id][SMARTTUB_CONTROLLER] + controller = entry.runtime_data entities = [ SmartTubThermostat(controller.coordinator, spa) for spa in controller.spas @@ -71,7 +72,9 @@ class SmartTubThermostat(SmartTubEntity, ClimateEntity): _attr_temperature_unit = UnitOfTemperature.CELSIUS _attr_preset_modes = list(PRESET_MODES.values()) - def __init__(self, coordinator, spa): + def __init__( + self, coordinator: DataUpdateCoordinator[dict[str, Any]], spa: Spa + ) -> None: """Initialize the entity.""" super().__init__(coordinator, spa, "Thermostat") diff --git a/homeassistant/components/smarttub/const.py b/homeassistant/components/smarttub/const.py index f97ef65a54c..dadc66da942 100644 --- a/homeassistant/components/smarttub/const.py +++ b/homeassistant/components/smarttub/const.py @@ -4,8 +4,6 @@ DOMAIN = "smarttub" EVENT_SMARTTUB = "smarttub" -SMARTTUB_CONTROLLER = "smarttub_controller" - SCAN_INTERVAL = 60 POLLING_TIMEOUT = 10 diff --git a/homeassistant/components/smarttub/controller.py b/homeassistant/components/smarttub/controller.py index 353e2093997..d8299bbd786 100644 --- a/homeassistant/components/smarttub/controller.py +++ b/homeassistant/components/smarttub/controller.py @@ -3,13 +3,15 @@ import asyncio from datetime import timedelta import logging +from typing import Any from aiohttp import client_exceptions -from smarttub import APIError, LoginFailed, SmartTub +from smarttub import APIError, LoginFailed, SmartTub, Spa from smarttub.api import Account +from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_EMAIL, CONF_PASSWORD -from homeassistant.core import callback +from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from homeassistant.helpers import device_registry as dr from homeassistant.helpers.aiohttp_client import async_get_clientsession @@ -29,19 +31,21 @@ from .helpers import get_spa_name _LOGGER = logging.getLogger(__name__) +type SmartTubConfigEntry = ConfigEntry[SmartTubController] + class SmartTubController: """Interface between Home Assistant and the SmartTub API.""" - def __init__(self, hass): + coordinator: DataUpdateCoordinator[dict[str, Any]] + spas: list[Spa] + _account: Account + + def __init__(self, hass: HomeAssistant) -> None: """Initialize an interface to SmartTub.""" self._hass = hass - self._account = None - self.spas = set() - self.coordinator = None - - async def async_setup_entry(self, entry): + async def async_setup_entry(self, entry: SmartTubConfigEntry) -> bool: """Perform initial setup. Authenticate, query static state, set up polling, and otherwise make @@ -79,7 +83,7 @@ class SmartTubController: return True - async def async_update_data(self): + async def async_update_data(self) -> dict[str, Any]: """Query the API and return the new state.""" data = {} @@ -92,7 +96,7 @@ class SmartTubController: return data - async def _get_spa_data(self, spa): + async def _get_spa_data(self, spa: Spa) -> dict[str, Any]: full_status, reminders, errors = await asyncio.gather( spa.get_status_full(), spa.get_reminders(), @@ -107,7 +111,7 @@ class SmartTubController: } @callback - def async_register_devices(self, entry): + def async_register_devices(self, entry: SmartTubConfigEntry) -> None: """Register devices with the device registry for all spas.""" device_registry = dr.async_get(self._hass) for spa in self.spas: @@ -119,11 +123,8 @@ class SmartTubController: model=spa.model, ) - async def login(self, email, password) -> Account: - """Retrieve the account corresponding to the specified email and password. - - Returns None if the credentials are invalid. - """ + async def login(self, email: str, password: str) -> Account: + """Retrieve the account corresponding to the specified email and password.""" api = SmartTub(async_get_clientsession(self._hass)) diff --git a/homeassistant/components/smarttub/entity.py b/homeassistant/components/smarttub/entity.py index f9ab1d10bfe..069fd50c5f2 100644 --- a/homeassistant/components/smarttub/entity.py +++ b/homeassistant/components/smarttub/entity.py @@ -1,6 +1,8 @@ """Base classes for SmartTub entities.""" -import smarttub +from typing import Any + +from smarttub import Spa, SpaState from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.update_coordinator import ( @@ -16,7 +18,10 @@ class SmartTubEntity(CoordinatorEntity): """Base class for SmartTub entities.""" def __init__( - self, coordinator: DataUpdateCoordinator, spa: smarttub.Spa, entity_name + self, + coordinator: DataUpdateCoordinator[dict[str, Any]], + spa: Spa, + entity_name: str, ) -> None: """Initialize the entity. @@ -36,7 +41,7 @@ class SmartTubEntity(CoordinatorEntity): self._attr_name = f"{spa_name} {entity_name}" @property - def spa_status(self) -> smarttub.SpaState: + def spa_status(self) -> SpaState: """Retrieve the result of Spa.get_status().""" return self.coordinator.data[self.spa.id].get("status") @@ -45,7 +50,13 @@ class SmartTubEntity(CoordinatorEntity): class SmartTubSensorBase(SmartTubEntity): """Base class for SmartTub sensors.""" - def __init__(self, coordinator, spa, sensor_name, state_key): + def __init__( + self, + coordinator: DataUpdateCoordinator[dict[str, Any]], + spa: Spa, + sensor_name: str, + state_key: str, + ) -> None: """Initialize the entity.""" super().__init__(coordinator, spa, sensor_name) self._state_key = state_key diff --git a/homeassistant/components/smarttub/light.py b/homeassistant/components/smarttub/light.py index dda936aa56a..b6e056d37e0 100644 --- a/homeassistant/components/smarttub/light.py +++ b/homeassistant/components/smarttub/light.py @@ -12,29 +12,24 @@ from homeassistant.components.light import ( LightEntity, LightEntityFeature, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from .const import ( - ATTR_LIGHTS, - DEFAULT_LIGHT_BRIGHTNESS, - DEFAULT_LIGHT_EFFECT, - DOMAIN, - SMARTTUB_CONTROLLER, -) +from .const import ATTR_LIGHTS, DEFAULT_LIGHT_BRIGHTNESS, DEFAULT_LIGHT_EFFECT +from .controller import SmartTubConfigEntry from .entity import SmartTubEntity from .helpers import get_spa_name async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: SmartTubConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up entities for any lights in the tub.""" - controller = hass.data[DOMAIN][entry.entry_id][SMARTTUB_CONTROLLER] + controller = entry.runtime_data entities = [ SmartTubLight(controller.coordinator, light) @@ -52,7 +47,9 @@ class SmartTubLight(SmartTubEntity, LightEntity): _attr_supported_color_modes = {ColorMode.BRIGHTNESS} _attr_supported_features = LightEntityFeature.EFFECT - def __init__(self, coordinator, light): + def __init__( + self, coordinator: DataUpdateCoordinator[dict[str, Any]], light: SpaLight + ) -> None: """Initialize the entity.""" super().__init__(coordinator, light.spa, "light") self.light_zone = light.zone diff --git a/homeassistant/components/smarttub/sensor.py b/homeassistant/components/smarttub/sensor.py index b2bb1170d09..5116bfb3aee 100644 --- a/homeassistant/components/smarttub/sensor.py +++ b/homeassistant/components/smarttub/sensor.py @@ -1,18 +1,19 @@ """Platform for sensor integration.""" from enum import Enum +from typing import Any import smarttub import voluptuous as vol from homeassistant.components.sensor import SensorEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers import config_validation as cv, entity_platform from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.typing import VolDictType +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from .const import DOMAIN, SMARTTUB_CONTROLLER +from .controller import SmartTubConfigEntry from .entity import SmartTubSensorBase # the desired duration, in hours, of the cycle @@ -44,12 +45,12 @@ SET_SECONDARY_FILTRATION_SCHEMA: VolDictType = { async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: SmartTubConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up sensor entities for the sensors in the tub.""" - controller = hass.data[DOMAIN][entry.entry_id][SMARTTUB_CONTROLLER] + controller = entry.runtime_data entities = [] for spa in controller.spas: @@ -107,7 +108,9 @@ class SmartTubSensor(SmartTubSensorBase, SensorEntity): class SmartTubPrimaryFiltrationCycle(SmartTubSensor): """The primary filtration cycle.""" - def __init__(self, coordinator, spa): + def __init__( + self, coordinator: DataUpdateCoordinator[dict[str, Any]], spa: smarttub.Spa + ) -> None: """Initialize the entity.""" super().__init__( coordinator, spa, "Primary Filtration Cycle", "primary_filtration" @@ -124,7 +127,7 @@ class SmartTubPrimaryFiltrationCycle(SmartTubSensor): return self.cycle.status.name.lower() @property - def extra_state_attributes(self): + def extra_state_attributes(self) -> dict[str, Any]: """Return the state attributes.""" return { ATTR_DURATION: self.cycle.duration, @@ -145,7 +148,9 @@ class SmartTubPrimaryFiltrationCycle(SmartTubSensor): class SmartTubSecondaryFiltrationCycle(SmartTubSensor): """The secondary filtration cycle.""" - def __init__(self, coordinator, spa): + def __init__( + self, coordinator: DataUpdateCoordinator[dict[str, Any]], spa: smarttub.Spa + ) -> None: """Initialize the entity.""" super().__init__( coordinator, spa, "Secondary Filtration Cycle", "secondary_filtration" @@ -162,7 +167,7 @@ class SmartTubSecondaryFiltrationCycle(SmartTubSensor): return self.cycle.status.name.lower() @property - def extra_state_attributes(self): + def extra_state_attributes(self) -> dict[str, Any]: """Return the state attributes.""" return { ATTR_CYCLE_LAST_UPDATED: self.cycle.last_updated.isoformat(), diff --git a/homeassistant/components/smarttub/switch.py b/homeassistant/components/smarttub/switch.py index 2dedad8e18a..12d15d63f9b 100644 --- a/homeassistant/components/smarttub/switch.py +++ b/homeassistant/components/smarttub/switch.py @@ -6,23 +6,24 @@ from typing import Any from smarttub import SpaPump from homeassistant.components.switch import SwitchEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from .const import API_TIMEOUT, ATTR_PUMPS, DOMAIN, SMARTTUB_CONTROLLER +from .const import API_TIMEOUT, ATTR_PUMPS +from .controller import SmartTubConfigEntry from .entity import SmartTubEntity from .helpers import get_spa_name async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: SmartTubConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up switch entities for the pumps on the tub.""" - controller = hass.data[DOMAIN][entry.entry_id][SMARTTUB_CONTROLLER] + controller = entry.runtime_data entities = [ SmartTubPump(controller.coordinator, pump) @@ -36,7 +37,9 @@ async def async_setup_entry( class SmartTubPump(SmartTubEntity, SwitchEntity): """A pump on a spa.""" - def __init__(self, coordinator, pump: SpaPump) -> None: + def __init__( + self, coordinator: DataUpdateCoordinator[dict[str, Any]], pump: SpaPump + ) -> None: """Initialize the entity.""" super().__init__(coordinator, pump.spa, "pump") self.pump_id = pump.id diff --git a/tests/components/smarttub/test_init.py b/tests/components/smarttub/test_init.py index b1eac3fd98b..ff27820fca1 100644 --- a/tests/components/smarttub/test_init.py +++ b/tests/components/smarttub/test_init.py @@ -4,7 +4,6 @@ from unittest.mock import patch from smarttub import LoginFailed -from homeassistant.components import smarttub from homeassistant.components.smarttub.const import DOMAIN from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState from homeassistant.core import HomeAssistant @@ -61,13 +60,13 @@ async def test_config_passed_to_config_entry( ) -> None: """Test that configured options are loaded via config entry.""" config_entry.add_to_hass(hass) - assert await async_setup_component(hass, smarttub.DOMAIN, config_data) + assert await async_setup_component(hass, DOMAIN, config_data) async def test_unload_entry(hass: HomeAssistant, config_entry) -> None: """Test being able to unload an entry.""" config_entry.add_to_hass(hass) - assert await async_setup_component(hass, smarttub.DOMAIN, {}) is True + assert await async_setup_component(hass, DOMAIN, {}) is True assert await hass.config_entries.async_unload(config_entry.entry_id)