Add restore logic to Shelly climate platform (#61632)

* Add restore logic to Shelly climate platform

* Handle missing channel on restore

homeassistant/components/shelly/__init__.py

@@ -69,7 +69,6 @@ from .utils import (
 BLOCK_PLATFORMS: Final = [
     Platform.BINARY_SENSOR,
     Platform.BUTTON,
-    Platform.CLIMATE,
     Platform.COVER,
     Platform.LIGHT,
     Platform.SENSOR,
@@ -77,6 +76,7 @@ BLOCK_PLATFORMS: Final = [
 ]
 BLOCK_SLEEPING_PLATFORMS: Final = [
     Platform.BINARY_SENSOR,
+    Platform.CLIMATE,
     Platform.SENSOR,
 ]
 RPC_PLATFORMS: Final = [

homeassistant/components/shelly/climate.py

@@ -3,12 +3,13 @@ from __future__ import annotations
 
 import asyncio
 import logging
+from types import MappingProxyType
 from typing import Any, Final, cast
 
 from aioshelly.block_device import Block
 import async_timeout
 
-from homeassistant.components.climate import ClimateEntity
+from homeassistant.components.climate import DOMAIN as CLIMATE_DOMAIN, ClimateEntity
 from homeassistant.components.climate.const import (
     CURRENT_HVAC_HEAT,
     CURRENT_HVAC_IDLE,
@@ -21,7 +22,9 @@ from homeassistant.components.climate.const import (
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TEMPERATURE, TEMP_CELSIUS
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, State, callback
+from homeassistant.helpers import device_registry, entity, entity_registry
+from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.restore_state import RestoreEntity
 
@@ -33,7 +36,6 @@ from .const import (
     DOMAIN,
     SHTRV_01_TEMPERATURE_SETTINGS,
 )
-from .entity import ShellyBlockEntity
 from .utils import get_device_entry_gen
 
 _LOGGER: Final = logging.getLogger(__name__)
@@ -49,10 +51,30 @@ async def async_setup_entry(
     if get_device_entry_gen(config_entry) == 2:
         return
 
+    wrapper: BlockDeviceWrapper = hass.data[DOMAIN][DATA_CONFIG_ENTRY][
+        config_entry.entry_id
+    ][BLOCK]
+
+    if wrapper.device.initialized:
+        await async_setup_climate_entities(async_add_entities, wrapper)
+    else:
+        await async_restore_climate_entities(
+            hass, config_entry, async_add_entities, wrapper
+        )
+
+
+async def async_setup_climate_entities(
+    async_add_entities: AddEntitiesCallback,
+    wrapper: BlockDeviceWrapper,
+) -> None:
+    """Set up online climate devices."""
+
+    _LOGGER.info("Setup online climate device %s", wrapper.name)
     device_block: Block | None = None
     sensor_block: Block | None = None
 
-    wrapper = hass.data[DOMAIN][DATA_CONFIG_ENTRY][config_entry.entry_id][BLOCK]
+    assert wrapper.device.blocks
+
     for block in wrapper.device.blocks:
         if block.type == "device":
             device_block = block
@@ -60,10 +82,37 @@ async def async_setup_entry(
             sensor_block = block
 
     if sensor_block and device_block:
-        async_add_entities([ShellyClimate(wrapper, sensor_block, device_block)])
+        async_add_entities([BlockSleepingClimate(wrapper, sensor_block, device_block)])
 
 
-class ShellyClimate(ShellyBlockEntity, RestoreEntity, ClimateEntity):
+async def async_restore_climate_entities(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+    wrapper: BlockDeviceWrapper,
+) -> None:
+    """Restore sleeping climate devices."""
+    _LOGGER.info("Setup sleeping climate device %s", wrapper.name)
+
+    ent_reg = await entity_registry.async_get_registry(hass)
+    entries = entity_registry.async_entries_for_config_entry(
+        ent_reg, config_entry.entry_id
+    )
+
+    for entry in entries:
+
+        if entry.domain != CLIMATE_DOMAIN:
+            continue
+
+        _LOGGER.debug("Found entry %s [%s]", entry.original_name, entry.domain)
+        async_add_entities([BlockSleepingClimate(wrapper, None, None, entry)])
+
+
+class BlockSleepingClimate(
+    RestoreEntity,
+    ClimateEntity,
+    entity.Entity,
+):
     """Representation of a Shelly climate device."""
 
     _attr_hvac_modes = [HVAC_MODE_OFF, HVAC_MODE_HEAT]
@@ -74,45 +123,77 @@ class ShellyClimate(ShellyBlockEntity, RestoreEntity, ClimateEntity):
     _attr_target_temperature_step = SHTRV_01_TEMPERATURE_SETTINGS["step"]
     _attr_temperature_unit = TEMP_CELSIUS
 
+    # pylint: disable=super-init-not-called
     def __init__(
-        self, wrapper: BlockDeviceWrapper, sensor_block: Block, device_block: Block
+        self,
+        wrapper: BlockDeviceWrapper,
+        sensor_block: Block | None,
+        device_block: Block | None,
+        entry: entity_registry.RegistryEntry | None = None,
     ) -> None:
         """Initialize climate."""
-        super().__init__(wrapper, sensor_block)
-
-        self.device_block = device_block
-
-        assert self.block.channel
 
+        self.wrapper = wrapper
+        self.block: Block | None = sensor_block
         self.control_result: dict[str, Any] | None = None
+        self.device_block: Block | None = device_block
+        self.last_state: State | None = None
+        self.last_state_attributes: MappingProxyType[str, Any]
+        self._preset_modes: list[str] = []
 
-        self._attr_name = self.wrapper.name
+        if self.block is not None and self.device_block is not None:
-        self._attr_unique_id = self.wrapper.mac
+            self._unique_id = f"{self.wrapper.mac}-{self.block.description}"
-        self._attr_preset_modes: list[str] = [
+            assert self.block.channel
-            PRESET_NONE,
+            self._preset_modes = [
-            *wrapper.device.settings["thermostats"][int(self.block.channel)][
+                PRESET_NONE,
-                "schedule_profile_names"
+                *wrapper.device.settings["thermostats"][int(self.block.channel)][
-            ],
+                    "schedule_profile_names"
-        ]
+                ],
+            ]
+        elif entry is not None:
+            self._unique_id = entry.unique_id
+
+    @property
+    def unique_id(self) -> str:
+        """Set unique id of entity."""
+        return self._unique_id
+
+    @property
+    def name(self) -> str:
+        """Name of entity."""
+        return self.wrapper.name
+
+    @property
+    def should_poll(self) -> bool:
+        """If device should be polled."""
+        return False
 
     @property
     def target_temperature(self) -> float | None:
         """Set target temperature."""
-        return cast(float, self.block.targetTemp)
+        if self.block is not None:
+            return cast(float, self.block.targetTemp)
+        return self.last_state_attributes.get("temperature")
 
     @property
     def current_temperature(self) -> float | None:
         """Return current temperature."""
-        return cast(float, self.block.temp)
+        if self.block is not None:
+            return cast(float, self.block.temp)
+        return self.last_state_attributes.get("current_temperature")
 
     @property
     def available(self) -> bool:
         """Device availability."""
-        return not cast(bool, self.device_block.valveError)
+        if self.device_block is not None:
+            return not cast(bool, self.device_block.valveError)
+        return True
 
     @property
     def hvac_mode(self) -> str:
         """HVAC current mode."""
+        if self.device_block is None:
+            return self.last_state.state if self.last_state else HVAC_MODE_OFF
         if self.device_block.mode is None or self._check_is_off():
             return HVAC_MODE_OFF
 
@@ -121,20 +202,45 @@ class ShellyClimate(ShellyBlockEntity, RestoreEntity, ClimateEntity):
     @property
     def preset_mode(self) -> str | None:
         """Preset current mode."""
+        if self.device_block is None:
+            return self.last_state_attributes.get("preset_mode")
         if self.device_block.mode is None:
-            return None
+            return PRESET_NONE
-        return self._attr_preset_modes[cast(int, self.device_block.mode)]
+        return self._preset_modes[cast(int, self.device_block.mode)]
 
     @property
     def hvac_action(self) -> str | None:
         """HVAC current action."""
-        if self.device_block.status is None or self._check_is_off():
+        if (
+            self.device_block is None
+            or self.device_block.status is None
+            or self._check_is_off()
+        ):
             return CURRENT_HVAC_OFF
 
         return (
             CURRENT_HVAC_IDLE if self.device_block.status == "0" else CURRENT_HVAC_HEAT
         )
 
+    @property
+    def preset_modes(self) -> list[str]:
+        """Preset available modes."""
+        return self._preset_modes
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Device info."""
+        return {
+            "connections": {(device_registry.CONNECTION_NETWORK_MAC, self.wrapper.mac)}
+        }
+
+    @property
+    def channel(self) -> str | None:
+        """Device channel."""
+        if self.block is not None:
+            return self.block.channel
+        return self.last_state_attributes.get("channel")
+
     def _check_is_off(self) -> bool:
         """Return if valve is off or on."""
         return bool(
@@ -148,7 +254,7 @@ class ShellyClimate(ShellyBlockEntity, RestoreEntity, ClimateEntity):
         try:
             async with async_timeout.timeout(AIOSHELLY_DEVICE_TIMEOUT_SEC):
                 return await self.wrapper.device.http_request(
-                    "get", f"thermostat/{self.block.channel}", kwargs
+                    "get", f"thermostat/{self.channel}", kwargs
                 )
         except (asyncio.TimeoutError, OSError) as err:
             _LOGGER.error(
@@ -186,3 +292,41 @@ class ShellyClimate(ShellyBlockEntity, RestoreEntity, ClimateEntity):
             await self.set_state_full_path(
                 schedule=1, schedule_profile=f"{preset_index}"
             )
+
+    async def async_added_to_hass(self) -> None:
+        """Handle entity which will be added."""
+        _LOGGER.info("Restoring entity %s", self.name)
+
+        last_state = await self.async_get_last_state()
+
+        if last_state is not None:
+            self.last_state = last_state
+            self.last_state_attributes = self.last_state.attributes
+            self._preset_modes = cast(
+                list, self.last_state.attributes.get("preset_modes")
+            )
+
+        self.async_on_remove(self.wrapper.async_add_listener(self._update_callback))
+
+    async def async_update(self) -> None:
+        """Update entity with latest info."""
+        await self.wrapper.async_request_refresh()
+
+    @callback
+    def _update_callback(self) -> None:
+        """Handle device update."""
+        if not self.wrapper.device.initialized:
+            self.async_write_ha_state()
+            return
+
+        assert self.wrapper.device.blocks
+
+        for block in self.wrapper.device.blocks:
+            if block.type == "device":
+                self.device_block = block
+            if hasattr(block, "targetTemp"):
+                self.block = block
+
+                _LOGGER.debug("Entity %s attached to block", self.name)
+                self.async_write_ha_state()
+                return