Use first known thermostat state as main state (pick #1054) (#1057)

Pick commit a044063526 from #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 home-assistant/core#121335

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 (

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