Bugfix: implement RestoreState and bump backend for Plugwise climate (home-assistant#155126)

Author: bouwew

homeassistant/components/plugwise/climate.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import Any
 
 from homeassistant.components.climate import (
@@ -13,18 +14,44 @@
     HVACAction,
     HVACMode,
 )
-from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature
+from homeassistant.const import ATTR_TEMPERATURE, STATE_OFF, STATE_ON, UnitOfTemperature
 from homeassistant.core import HomeAssistant, callback
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+from homeassistant.helpers.restore_state import ExtraStoredData, RestoreEntity
 
 from .const import DOMAIN, MASTER_THERMOSTATS
 from .coordinator import PlugwiseConfigEntry, PlugwiseDataUpdateCoordinator
 from .entity import PlugwiseEntity
 from .util import plugwise_command
 
+ERROR_NO_SCHEDULE = "set_schedule_first"
 PARALLEL_UPDATES = 0
 
 
+@dataclass
+class PlugwiseClimateExtraStoredData(ExtraStoredData):
+    """Object to hold extra stored data."""
+
+    last_active_schedule: str | None
+    previous_action_mode: str | None
+
+    def as_dict(self) -> dict[str, Any]:
+        """Return a dict representation of the text data."""
+        return {
+            "last_active_schedule": self.last_active_schedule,
+            "previous_action_mode": self.previous_action_mode,
+        }
+
+    @classmethod
+    def from_dict(cls, restored: dict[str, Any]) -> PlugwiseClimateExtraStoredData:
+        """Initialize a stored data object from a dict."""
+        return cls(
+            last_active_schedule=restored.get("last_active_schedule"),
+            previous_action_mode=restored.get("previous_action_mode"),
+        )
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     entry: PlugwiseConfigEntry,
@@ -56,14 +83,26 @@ def _add_entities() -> None:
     entry.async_on_unload(coordinator.async_add_listener(_add_entities))
 
 
-class PlugwiseClimateEntity(PlugwiseEntity, ClimateEntity):
+class PlugwiseClimateEntity(PlugwiseEntity, ClimateEntity, RestoreEntity):
     """Representation of a Plugwise thermostat."""
 
     _attr_name = None
     _attr_temperature_unit = UnitOfTemperature.CELSIUS
     _attr_translation_key = DOMAIN
 
-    _previous_mode: str = "heating"
+    _last_active_schedule: str | None = None
+    _previous_action_mode: str | None = HVACAction.HEATING.value
+
+    async def async_added_to_hass(self) -> None:
+        """Run when entity about to be added."""
+        await super().async_added_to_hass()
+
+        if extra_data := await self.async_get_last_extra_data():
+            plugwise_extra_data = PlugwiseClimateExtraStoredData.from_dict(
+                extra_data.as_dict()
+            )
+            self._last_active_schedule = plugwise_extra_data.last_active_schedule
+            self._previous_action_mode = plugwise_extra_data.previous_action_mode
 
     def __init__(
         self,
@@ -76,7 +115,6 @@ def __init__(
 
         gateway_id: str = coordinator.api.gateway_id
         self._gateway_data = coordinator.data[gateway_id]
-
         self._location = device_id
         if (location := self.device.get("location")) is not None:
             self._location = location
@@ -105,25 +143,19 @@ def __init__(
             self.device["thermostat"]["resolution"], 0.1
         )
 
-    def _previous_action_mode(self, coordinator: PlugwiseDataUpdateCoordinator) -> None:
-        """Return the previous action-mode when the regulation-mode is not heating or cooling.
-
-        Helper for set_hvac_mode().
-        """
-        # When no cooling available, _previous_mode is always heating
-        if (
-            "regulation_modes" in self._gateway_data
-            and "cooling" in self._gateway_data["regulation_modes"]
-        ):
-            mode = self._gateway_data["select_regulation_mode"]
-            if mode in ("cooling", "heating"):
-                self._previous_mode = mode
-
     @property
     def current_temperature(self) -> float:
         """Return the current temperature."""
         return self.device["sensors"]["temperature"]
 
+    @property
+    def extra_restore_state_data(self) -> PlugwiseClimateExtraStoredData:
+        """Return text specific state data to be restored."""
+        return PlugwiseClimateExtraStoredData(
+            last_active_schedule=self._last_active_schedule,
+            previous_action_mode=self._previous_action_mode,
+        )
+
     @property
     def target_temperature(self) -> float:
         """Return the temperature we try to reach.
@@ -170,9 +202,10 @@ def hvac_modes(self) -> list[HVACMode]:
 
         if self.coordinator.api.cooling_present:
             if "regulation_modes" in self._gateway_data:
-                if self._gateway_data["select_regulation_mode"] == "cooling":
+                selected = self._gateway_data.get("select_regulation_mode")
+                if selected == HVACAction.COOLING.value:
                     hvac_modes.append(HVACMode.COOL)
-                if self._gateway_data["select_regulation_mode"] == "heating":
+                if selected == HVACAction.HEATING.value:
                     hvac_modes.append(HVACMode.HEAT)
             else:
                 hvac_modes.append(HVACMode.HEAT_COOL)
@@ -184,8 +217,16 @@ def hvac_modes(self) -> list[HVACMode]:
     @property
     def hvac_action(self) -> HVACAction:
         """Return the current running hvac operation if supported."""
-        # Keep track of the previous action-mode
-        self._previous_action_mode(self.coordinator)
+        # Keep track of the previous hvac_action mode.
+        # When no cooling available, _previous_action_mode is always heating
+        if (
+            "regulation_modes" in self._gateway_data
+            and HVACAction.COOLING.value in self._gateway_data["regulation_modes"]
+        ):
+            mode = self._gateway_data["select_regulation_mode"]
+            if mode in (HVACAction.COOLING.value, HVACAction.HEATING.value):
+                self._previous_action_mode = mode
+
         if (action := self.device.get("control_state")) is not None:
             return HVACAction(action)
 
@@ -219,14 +260,33 @@ async def async_set_hvac_mode(self, hvac_mode: HVACMode) -> None:
             return
 
         if hvac_mode == HVACMode.OFF:
-            await self.coordinator.api.set_regulation_mode(hvac_mode)
+            await self.coordinator.api.set_regulation_mode(hvac_mode.value)
         else:
+            current = self.device.get("select_schedule")
+            desired = current
+
+            # Capture the last valid schedule
+            if desired and desired != "off":
+                self._last_active_schedule = desired
+            elif desired == "off":
+                desired = self._last_active_schedule
+
+            # Enabling HVACMode.AUTO requires a previously set schedule for saving and restoring
+            if hvac_mode == HVACMode.AUTO and not desired:
+                raise HomeAssistantError(
+                    translation_domain=DOMAIN,
+                    translation_key=ERROR_NO_SCHEDULE,
+                )
+
             await self.coordinator.api.set_schedule_state(
                 self._location,
-                "on" if hvac_mode == HVACMode.AUTO else "off",
+                STATE_ON if hvac_mode == HVACMode.AUTO else STATE_OFF,
+                desired,
             )
-            if self.hvac_mode == HVACMode.OFF:
-                await self.coordinator.api.set_regulation_mode(self._previous_mode)
+            if self.hvac_mode == HVACMode.OFF and self._previous_action_mode:
+                await self.coordinator.api.set_regulation_mode(
+                    self._previous_action_mode
+                )
 
     @plugwise_command
     async def async_set_preset_mode(self, preset_mode: str) -> None:

homeassistant/components/plugwise/manifest.json

@@ -8,6 +8,6 @@
   "iot_class": "local_polling",
   "loggers": ["plugwise"],
   "quality_scale": "platinum",
-  "requirements": ["plugwise==1.8.2"],
+  "requirements": ["plugwise==1.8.3"],
   "zeroconf": ["_plugwise._tcp.local."]
 }

homeassistant/components/plugwise/strings.json

@@ -314,6 +314,9 @@
     "invalid_xml_data": {
       "message": "[%key:component::plugwise::config::error::response_error%]"
     },
+    "set_schedule_first": {
+      "message": "Failed setting HVACMode, set a schedule first."
+    },
     "unsupported_firmware": {
       "message": "[%key:component::plugwise::config::error::unsupported%]"
     }

requirements_all.txt

@@ -23,12 +23,18 @@
     HVACAction,
     HVACMode,
 )
-from homeassistant.const import ATTR_ENTITY_ID, ATTR_TEMPERATURE
-from homeassistant.core import HomeAssistant
+from homeassistant.components.plugwise.climate import PlugwiseClimateExtraStoredData
+from homeassistant.const import ATTR_ENTITY_ID, ATTR_TEMPERATURE, STATE_OFF, STATE_ON
+from homeassistant.core import HomeAssistant, State
 from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
 from homeassistant.helpers import entity_registry as er
 
-from tests.common import MockConfigEntry, async_fire_time_changed, snapshot_platform
+from tests.common import (
+    MockConfigEntry,
+    async_fire_time_changed,
+    mock_restore_cache_with_extra_data,
+    snapshot_platform,
+)
 
 HA_PLUGWISE_SMILE_ASYNC_UPDATE = (
     "homeassistant.components.plugwise.coordinator.Smile.async_update"
@@ -105,7 +111,9 @@ async def test_adam_climate_entity_climate_changes(
     )
     assert mock_smile_adam.set_schedule_state.call_count == 2
     mock_smile_adam.set_schedule_state.assert_called_with(
-        "c50f167537524366a5af7aa3942feb1e", HVACMode.OFF
+        "c50f167537524366a5af7aa3942feb1e",
+        STATE_OFF,
+        "GF7  Woonkamer",
     )
 
     with pytest.raises(
@@ -138,6 +146,98 @@ async def test_adam_climate_adjust_negative_testing(
         )
 
 
+@pytest.mark.parametrize("chosen_env", ["m_adam_heating"], indirect=True)
+@pytest.mark.parametrize("cooling_present", [False], indirect=True)
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+async def test_adam_restore_state_climate(
+    hass: HomeAssistant,
+    mock_smile_adam_heat_cool: MagicMock,
+    mock_config_entry: MockConfigEntry,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test restore_state for climate with restored schedule."""
+    mock_restore_cache_with_extra_data(
+        hass,
+        [
+            (
+                State("climate.living_room", "heat"),
+                PlugwiseClimateExtraStoredData(
+                    last_active_schedule=None,
+                    previous_action_mode="heating",
+                ).as_dict(),
+            ),
+            (
+                State("climate.bathroom", "heat"),
+                PlugwiseClimateExtraStoredData(
+                    last_active_schedule="Badkamer",
+                    previous_action_mode=None,
+                ).as_dict(),
+            ),
+        ],
+    )
+
+    mock_config_entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert (state := hass.states.get("climate.living_room"))
+    assert state.state == "heat"
+
+    # Verify a HomeAssistantError is raised setting a schedule with last_active_schedule = None
+    with pytest.raises(HomeAssistantError):
+        await hass.services.async_call(
+            CLIMATE_DOMAIN,
+            SERVICE_SET_HVAC_MODE,
+            {ATTR_ENTITY_ID: "climate.living_room", ATTR_HVAC_MODE: HVACMode.AUTO},
+            blocking=True,
+        )
+
+    data = mock_smile_adam_heat_cool.async_update.return_value
+    data["f2bf9048bef64cc5b6d5110154e33c81"]["climate_mode"] = "off"
+    data["da224107914542988a88561b4452b0f6"]["selec_regulation_mode"] = "off"
+    with patch(HA_PLUGWISE_SMILE_ASYNC_UPDATE, return_value=data):
+        freezer.tick(timedelta(minutes=1))
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
+
+        assert (state := hass.states.get("climate.living_room"))
+        assert state.state == "off"
+
+        # Verify restoration of previous_action_mode = heating
+        await hass.services.async_call(
+            CLIMATE_DOMAIN,
+            SERVICE_SET_HVAC_MODE,
+            {ATTR_ENTITY_ID: "climate.living_room", ATTR_HVAC_MODE: HVACMode.HEAT},
+            blocking=True,
+        )
+        # Verify set_schedule_state was called with the restored schedule
+        mock_smile_adam_heat_cool.set_regulation_mode.assert_called_with(
+            "heating",
+        )
+
+    data = mock_smile_adam_heat_cool.async_update.return_value
+    data["f871b8c4d63549319221e294e4f88074"]["climate_mode"] = "heat"
+    with patch(HA_PLUGWISE_SMILE_ASYNC_UPDATE, return_value=data):
+        freezer.tick(timedelta(minutes=1))
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
+
+        assert (state := hass.states.get("climate.bathroom"))
+        assert state.state == "heat"
+
+        # Verify restoration is used when setting a schedule
+        await hass.services.async_call(
+            CLIMATE_DOMAIN,
+            SERVICE_SET_HVAC_MODE,
+            {ATTR_ENTITY_ID: "climate.bathroom", ATTR_HVAC_MODE: HVACMode.AUTO},
+            blocking=True,
+        )
+        # Verify set_schedule_state was called with the restored schedule
+        mock_smile_adam_heat_cool.set_schedule_state.assert_called_with(
+            "f871b8c4d63549319221e294e4f88074", STATE_ON, "Badkamer"
+        )
+
+
 @pytest.mark.parametrize("chosen_env", ["m_adam_heating"], indirect=True)
 @pytest.mark.parametrize("cooling_present", [False], indirect=True)
 @pytest.mark.parametrize("platforms", [(CLIMATE_DOMAIN,)])
@@ -173,6 +273,7 @@ async def test_adam_3_climate_entity_attributes(
     ]
     data = mock_smile_adam_heat_cool.async_update.return_value
     data["da224107914542988a88561b4452b0f6"]["select_regulation_mode"] = "heating"
+    data["f2bf9048bef64cc5b6d5110154e33c81"]["climate_mode"] = "heat"
     data["f2bf9048bef64cc5b6d5110154e33c81"]["control_state"] = HVACAction.HEATING
     data["056ee145a816487eaa69243c3280f8bf"]["binary_sensors"]["cooling_state"] = False
     data["056ee145a816487eaa69243c3280f8bf"]["binary_sensors"]["heating_state"] = True
@@ -193,6 +294,7 @@ async def test_adam_3_climate_entity_attributes(
 
     data = mock_smile_adam_heat_cool.async_update.return_value
     data["da224107914542988a88561b4452b0f6"]["select_regulation_mode"] = "cooling"
+    data["f2bf9048bef64cc5b6d5110154e33c81"]["climate_mode"] = "cool"
     data["f2bf9048bef64cc5b6d5110154e33c81"]["control_state"] = HVACAction.COOLING
     data["056ee145a816487eaa69243c3280f8bf"]["binary_sensors"]["cooling_state"] = True
     data["056ee145a816487eaa69243c3280f8bf"]["binary_sensors"]["heating_state"] = False
@@ -334,7 +436,9 @@ async def test_anna_climate_entity_climate_changes(
     )
     assert mock_smile_anna.set_schedule_state.call_count == 1
     mock_smile_anna.set_schedule_state.assert_called_with(
-        "c784ee9fdab44e1395b8dee7d7a497d5", HVACMode.OFF
+        "c784ee9fdab44e1395b8dee7d7a497d5",
+        STATE_OFF,
+        "standaard",
     )
 
     # Mock user deleting last schedule from app or browser