From a0440635260abe2d8b6719ff2260510d2fed9b3f Mon Sep 17 00:00:00 2001
From: Teemu R <tpr@iki.fi>
Date: Thu, 11 Jul 2024 15:11:50 +0200
Subject: [PATCH] Use first known thermostat state as main state (#1054)

Instead of trying to use the first state when multiple are reported,
iterate over the known states and pick the first matching.
This will fix an issue where the device reports extra states (like
`low_battery`) while having a known mode active.

Related to https://github.com/home-assistant/core/issues/121335
---
 kasa/smart/modules/temperaturecontrol.py      | 43 ++++++++++---------
 .../smart/modules/test_temperaturecontrol.py  | 10 ++++-
 2 files changed, 31 insertions(+), 22 deletions(-)

diff --git a/kasa/smart/modules/temperaturecontrol.py b/kasa/smart/modules/temperaturecontrol.py
index dcd0da72..00afe5b5 100644
--- a/kasa/smart/modules/temperaturecontrol.py
+++ b/kasa/smart/modules/temperaturecontrol.py
@@ -4,15 +4,10 @@ from __future__ import annotations
 
 import logging
 from enum import Enum
-from typing import TYPE_CHECKING
 
 from ...feature import Feature
 from ..smartmodule import SmartModule
 
-if TYPE_CHECKING:
-    from ..smartdevice import SmartDevice
-
-
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -31,11 +26,11 @@ class TemperatureControl(SmartModule):
 
     REQUIRED_COMPONENT = "temp_control"
 
-    def __init__(self, device: SmartDevice, module: str):
-        super().__init__(device, module)
+    def _initialize_features(self):
+        """Initialize features after the initial update."""
         self._add_feature(
             Feature(
-                device,
+                self._device,
                 id="target_temperature",
                 name="Target temperature",
                 container=self,
@@ -50,7 +45,7 @@ class TemperatureControl(SmartModule):
         # TODO: this might belong into its own module, temperature_correction?
         self._add_feature(
             Feature(
-                device,
+                self._device,
                 id="temperature_offset",
                 name="Temperature offset",
                 container=self,
@@ -65,7 +60,7 @@ class TemperatureControl(SmartModule):
 
         self._add_feature(
             Feature(
-                device,
+                self._device,
                 id="state",
                 name="State",
                 container=self,
@@ -78,7 +73,7 @@ class TemperatureControl(SmartModule):
 
         self._add_feature(
             Feature(
-                device,
+                self._device,
                 id="thermostat_mode",
                 name="Thermostat mode",
                 container=self,
@@ -109,23 +104,24 @@ class TemperatureControl(SmartModule):
         if self._device.sys_info.get("frost_protection_on", False):
             return ThermostatState.Off
 
-        states = self._device.sys_info["trv_states"]
+        states = self.states
 
         # If the states is empty, the device is idling
         if not states:
             return ThermostatState.Idle
 
+        # Discard known extra states, and report on unknown extra states
+        states.discard("low_battery")
         if len(states) > 1:
-            _LOGGER.warning(
-                "Got multiple states (%s), using the first one: %s", states, states[0]
-            )
+            _LOGGER.warning("Got multiple states: %s", states)
 
-        state = states[0]
-        try:
-            return ThermostatState(state)
-        except:  # noqa: E722
-            _LOGGER.warning("Got unknown state: %s", state)
-            return ThermostatState.Unknown
+        # Return the first known state
+        for state in ThermostatState:
+            if state.value in states:
+                return state
+
+        _LOGGER.warning("Got unknown state: %s", states)
+        return ThermostatState.Unknown
 
     @property
     def allowed_temperature_range(self) -> tuple[int, int]:
@@ -147,6 +143,11 @@ class TemperatureControl(SmartModule):
         """Return target temperature."""
         return self._device.sys_info["target_temp"]
 
+    @property
+    def states(self) -> set:
+        """Return thermostat states."""
+        return set(self._device.sys_info["trv_states"])
+
     async def set_target_temperature(self, target: float):
         """Set target temperature."""
         if (
diff --git a/kasa/tests/smart/modules/test_temperaturecontrol.py b/kasa/tests/smart/modules/test_temperaturecontrol.py
index 16e01ed2..90f91216 100644
--- a/kasa/tests/smart/modules/test_temperaturecontrol.py
+++ b/kasa/tests/smart/modules/test_temperaturecontrol.py
@@ -94,7 +94,7 @@ async def test_temperature_offset(dev):
         ),
         pytest.param(
             ThermostatState.Heating,
-            [ThermostatState.Heating],
+            ["heating"],
             False,
             id="heating is heating",
         ),
@@ -135,3 +135,11 @@ async def test_thermostat_mode_warnings(dev, mode, states, msg, caplog):
     temp_module.data["trv_states"] = states
     assert temp_module.mode is mode
     assert msg in caplog.text
+
+
+@thermostats_smart
+async def test_thermostat_heating_with_low_battery(dev):
+    """Test that mode is reported correctly with extra states."""
+    temp_module: TemperatureControl = dev.modules["TemperatureControl"]
+    temp_module.data["trv_states"] = ["low_battery", "heating"]
+    assert temp_module.mode is ThermostatState.Heating