Use runtime_data in smarttub (#145279)

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

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 {}

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")
 

homeassistant/components/smarttub/const.py

@@ -4,8 +4,6 @@ DOMAIN = "smarttub"
 
 EVENT_SMARTTUB = "smarttub"
 
-SMARTTUB_CONTROLLER = "smarttub_controller"
-
 SCAN_INTERVAL = 60
 
 POLLING_TIMEOUT = 10

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: SmartTubConfigEntry) -> bool:
-
-    async def async_setup_entry(self, entry):
         """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:
+    async def login(self, email: str, password: str) -> Account:
-        """Retrieve the account corresponding to the specified email and password.
+        """Retrieve the account corresponding to the specified email and password."""
-
-        Returns None if the credentials are invalid.
-        """
 
         api = SmartTub(async_get_clientsession(self._hass))
 

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

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 (
+from .const import ATTR_LIGHTS, DEFAULT_LIGHT_BRIGHTNESS, DEFAULT_LIGHT_EFFECT
-    ATTR_LIGHTS,
+from .controller import SmartTubConfigEntry
-    DEFAULT_LIGHT_BRIGHTNESS,
-    DEFAULT_LIGHT_EFFECT,
-    DOMAIN,
-    SMARTTUB_CONTROLLER,
-)
 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

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(),

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

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)