Improve temperature controls (#872)

This improves the temperature control features to allow implementing
climate platform support for homeassistant.

Also adds frostprotection module, which is also used to turn
the thermostat on and off.

kasa/smart/modules/__init__.py

@@ -12,6 +12,7 @@ from .devicemodule import DeviceModule
 from .energymodule import EnergyModule
 from .fanmodule import FanModule
 from .firmware import Firmware
+from .frostprotection import FrostProtectionModule
 from .humidity import HumiditySensor
 from .ledmodule import LedModule
 from .lighttransitionmodule import LightTransitionModule
@@ -42,4 +43,5 @@ __all__ = [
     "ColorTemperatureModule",
     "ColorModule",
     "WaterleakSensor",
+    "FrostProtectionModule",
 ]

kasa/smart/modules/frostprotection.py

@@ -0,0 +1,58 @@
+"""Frost protection module."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from ...feature import Feature
+from ..smartmodule import SmartModule
+
+# TODO: this may not be necessary with __future__.annotations
+if TYPE_CHECKING:
+    from ..smartdevice import SmartDevice
+
+
+class FrostProtectionModule(SmartModule):
+    """Implementation for frost protection module.
+
+    This basically turns the thermostat on and off.
+    """
+
+    REQUIRED_COMPONENT = "frost_protection"
+    # TODO: the information required for current features do not require this query
+    QUERY_GETTER_NAME = "get_frost_protection"
+
+    def __init__(self, device: SmartDevice, module: str):
+        super().__init__(device, module)
+        self._add_feature(
+            Feature(
+                device,
+                name="Frost protection enabled",
+                container=self,
+                attribute_getter="enabled",
+                attribute_setter="set_enabled",
+                type=Feature.Type.Switch,
+            )
+        )
+
+    @property
+    def enabled(self) -> bool:
+        """Return True if frost protection is on."""
+        return self._device.sys_info["frost_protection_on"]
+
+    async def set_enabled(self, enable: bool):
+        """Enable/disable frost protection."""
+        return await self.call(
+            "set_device_info",
+            {"frost_protection_on": enable},
+        )
+
+    @property
+    def minimum_temperature(self) -> int:
+        """Return frost protection minimum temperature."""
+        return self.data["min_temp"]
+
+    @property
+    def temperature_unit(self) -> str:
+        """Return frost protection temperature unit."""
+        return self.data["temp_unit"]

kasa/smart/modules/humidity.py

@@ -26,6 +26,7 @@ class HumiditySensor(SmartModule):
                 container=self,
                 attribute_getter="humidity",
                 icon="mdi:water-percent",
+                unit="%",
             )
         )
         self._add_feature(

kasa/smart/modules/temperature.py

@@ -26,6 +26,7 @@ class TemperatureSensor(SmartModule):
                 container=self,
                 attribute_getter="temperature",
                 icon="mdi:thermometer",
+                category=Feature.Category.Primary,
             )
         )
         if "current_temp_exception" in device.sys_info:

kasa/smart/modules/temperaturecontrol.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import logging
+from enum import Enum
 from typing import TYPE_CHECKING
 
 from ...feature import Feature
@@ -11,6 +13,19 @@ if TYPE_CHECKING:
     from ..smartdevice import SmartDevice
 
 
+_LOGGER = logging.getLogger(__name__)
+
+
+class ThermostatState(Enum):
+    """Thermostat state."""
+
+    Heating = "heating"
+    Calibrating = "progress_calibration"
+    Idle = "idle"
+    Off = "off"
+    Unknown = "unknown"
+
+
 class TemperatureControl(SmartModule):
     """Implementation of temperature module."""
 
@@ -25,8 +40,10 @@ class TemperatureControl(SmartModule):
                 container=self,
                 attribute_getter="target_temperature",
                 attribute_setter="set_target_temperature",
+                range_getter="allowed_temperature_range",
                 icon="mdi:thermometer",
                 type=Feature.Type.Number,
+                category=Feature.Category.Primary,
             )
         )
         # TODO: this might belong into its own module, temperature_correction?
@@ -40,6 +57,29 @@ class TemperatureControl(SmartModule):
                 minimum_value=-10,
                 maximum_value=10,
                 type=Feature.Type.Number,
+                category=Feature.Category.Config,
+            )
+        )
+
+        self._add_feature(
+            Feature(
+                device,
+                "State",
+                container=self,
+                attribute_getter="state",
+                attribute_setter="set_state",
+                category=Feature.Category.Primary,
+                type=Feature.Type.Switch,
+            )
+        )
+
+        self._add_feature(
+            Feature(
+                device,
+                "Mode",
+                container=self,
+                attribute_getter="mode",
+                category=Feature.Category.Primary,
             )
         )
 
@@ -48,6 +88,45 @@ class TemperatureControl(SmartModule):
         # Target temperature is contained in the main device info response.
         return {}
 
+    @property
+    def state(self) -> bool:
+        """Return thermostat state."""
+        return self._device.sys_info["frost_protection_on"] is False
+
+    async def set_state(self, enabled: bool):
+        """Set thermostat state."""
+        return await self.call("set_device_info", {"frost_protection_on": not enabled})
+
+    @property
+    def mode(self) -> ThermostatState:
+        """Return thermostat state."""
+        # If frost protection is enabled, the thermostat is off.
+        if self._device.sys_info.get("frost_protection_on", False):
+            return ThermostatState.Off
+
+        states = self._device.sys_info["trv_states"]
+
+        # If the states is empty, the device is idling
+        if not states:
+            return ThermostatState.Idle
+
+        if len(states) > 1:
+            _LOGGER.warning(
+                "Got multiple states (%s), using the first one: %s", states, states[0]
+            )
+
+        state = states[0]
+        try:
+            return ThermostatState(state)
+        except:  # noqa: E722
+            _LOGGER.warning("Got unknown state: %s", state)
+            return ThermostatState.Unknown
+
+    @property
+    def allowed_temperature_range(self) -> tuple[int, int]:
+        """Return allowed temperature range."""
+        return self.minimum_target_temperature, self.maximum_target_temperature
+
     @property
     def minimum_target_temperature(self) -> int:
         """Minimum available target temperature."""
@@ -74,7 +153,12 @@ class TemperatureControl(SmartModule):
                 f"[{self.minimum_target_temperature},{self.maximum_target_temperature}]"
             )
 
-        return await self.call("set_device_info", {"target_temp": target})
+        payload = {"target_temp": target}
+        # If the device has frost protection, we set it off to enable heating
+        if "frost_protection_on" in self._device.sys_info:
+            payload["frost_protection_on"] = False
+
+        return await self.call("set_device_info", payload)
 
     @property
     def temperature_offset(self) -> int:

kasa/tests/smart/modules/test_temperaturecontrol.py

@@ -1,6 +1,9 @@
+import logging
+
 import pytest
 
-from kasa.smart.modules import TemperatureSensor
+from kasa.smart.modules import TemperatureControl
+from kasa.smart.modules.temperaturecontrol import ThermostatState
 from kasa.tests.device_fixtures import parametrize, thermostats_smart
 
 temperature = parametrize(
@@ -20,7 +23,7 @@ temperature = parametrize(
 )
 async def test_temperature_control_features(dev, feature, type):
     """Test that features are registered and work as expected."""
-    temp_module: TemperatureSensor = dev.modules["TemperatureControl"]
+    temp_module: TemperatureControl = dev.modules["TemperatureControl"]
 
     prop = getattr(temp_module, feature)
     assert isinstance(prop, type)
@@ -32,3 +35,103 @@ async def test_temperature_control_features(dev, feature, type):
     await feat.set_value(10)
     await dev.update()
     assert feat.value == 10
+
+
+@thermostats_smart
+async def test_set_temperature_turns_heating_on(dev):
+    """Test that set_temperature turns heating on."""
+    temp_module: TemperatureControl = dev.modules["TemperatureControl"]
+
+    await temp_module.set_state(False)
+    await dev.update()
+    assert temp_module.state is False
+    assert temp_module.mode is ThermostatState.Off
+
+    await temp_module.set_target_temperature(10)
+    await dev.update()
+    assert temp_module.state is True
+    assert temp_module.mode is ThermostatState.Heating
+    assert temp_module.target_temperature == 10
+
+
+@thermostats_smart
+async def test_set_temperature_invalid_values(dev):
+    """Test that out-of-bounds temperature values raise errors."""
+    temp_module: TemperatureControl = dev.modules["TemperatureControl"]
+
+    with pytest.raises(ValueError):
+        await temp_module.set_target_temperature(-1)
+
+    with pytest.raises(ValueError):
+        await temp_module.set_target_temperature(100)
+
+
+@thermostats_smart
+async def test_temperature_offset(dev):
+    """Test the temperature offset API."""
+    temp_module: TemperatureControl = dev.modules["TemperatureControl"]
+    with pytest.raises(ValueError):
+        await temp_module.set_temperature_offset(100)
+
+    with pytest.raises(ValueError):
+        await temp_module.set_temperature_offset(-100)
+
+    await temp_module.set_temperature_offset(5)
+    await dev.update()
+    assert temp_module.temperature_offset == 5
+
+
+@thermostats_smart
+@pytest.mark.parametrize(
+    "mode, states, frost_protection",
+    [
+        pytest.param(ThermostatState.Idle, [], False, id="idle has empty"),
+        pytest.param(
+            ThermostatState.Off,
+            ["anything"],
+            True,
+            id="any state with frost_protection on means off",
+        ),
+        pytest.param(
+            ThermostatState.Heating,
+            [ThermostatState.Heating],
+            False,
+            id="heating is heating",
+        ),
+        pytest.param(ThermostatState.Unknown, ["invalid"], False, id="unknown state"),
+    ],
+)
+async def test_thermostat_mode(dev, mode, states, frost_protection):
+    """Test different thermostat modes."""
+    temp_module: TemperatureControl = dev.modules["TemperatureControl"]
+
+    temp_module.data["frost_protection_on"] = frost_protection
+    temp_module.data["trv_states"] = states
+
+    assert temp_module.state is not frost_protection
+    assert temp_module.mode is mode
+
+
+@thermostats_smart
+@pytest.mark.parametrize(
+    "mode, states, msg",
+    [
+        pytest.param(
+            ThermostatState.Heating,
+            ["heating", "something else"],
+            "Got multiple states",
+            id="multiple states",
+        ),
+        pytest.param(
+            ThermostatState.Unknown, ["foobar"], "Got unknown state", id="unknown state"
+        ),
+    ],
+)
+async def test_thermostat_mode_warnings(dev, mode, states, msg, caplog):
+    """Test thermostat modes that should log a warning."""
+    temp_module: TemperatureControl = dev.modules["TemperatureControl"]
+    caplog.set_level(logging.WARNING)
+
+    temp_module.data["trv_states"] = states
+    assert temp_module.mode is mode
+    assert msg in caplog.text