From 23c5ee089a09e2fc4a433947ad1265a4e0f1cd9f Mon Sep 17 00:00:00 2001
From: Teemu R <tpr@iki.fi>
Date: Wed, 22 May 2024 16:52:00 +0200
Subject: [PATCH] Add state feature for iot devices (#924)

This is allows a generic implementation for the switch platform in the
homeassistant integration.

Also elevates set_state(bool) to be part of the standard API.
---
 kasa/device.py        |  8 ++++++++
 kasa/iot/iotdevice.py | 19 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/kasa/device.py b/kasa/device.py
index 7156a219..d462239d 100644
--- a/kasa/device.py
+++ b/kasa/device.py
@@ -138,6 +138,14 @@ class Device(ABC):
     async def turn_off(self, **kwargs) -> dict | None:
         """Turn off the device."""
 
+    @abstractmethod
+    async def set_state(self, on: bool):
+        """Set the device state to *on*.
+
+        This allows turning the device on and off.
+        See also *turn_off* and *turn_on*.
+        """
+
     @property
     def host(self) -> str:
         """The device host."""
diff --git a/kasa/iot/iotdevice.py b/kasa/iot/iotdevice.py
index 25e3b44d..dfe48a12 100755
--- a/kasa/iot/iotdevice.py
+++ b/kasa/iot/iotdevice.py
@@ -323,6 +323,18 @@ class IotDevice(Device):
         """Initialize modules not added in init."""
 
     async def _initialize_features(self):
+        """Initialize common features."""
+        self._add_feature(
+            Feature(
+                self,
+                id="state",
+                name="State",
+                attribute_getter="is_on",
+                attribute_setter="set_state",
+                type=Feature.Type.Switch,
+                category=Feature.Category.Primary,
+            )
+        )
         self._add_feature(
             Feature(
                 device=self,
@@ -634,6 +646,13 @@ class IotDevice(Device):
         """Return True if the device is on."""
         raise NotImplementedError("Device subclass needs to implement this.")
 
+    async def set_state(self, on: bool):
+        """Set the device state."""
+        if on:
+            return await self.turn_on()
+        else:
+            return await self.turn_off()
+
     @property  # type: ignore
     @requires_update
     def on_since(self) -> datetime | None: