From e1e2a396b8e14980f12ae5329403fc4cf9b23db7 Mon Sep 17 00:00:00 2001
From: Steven B <51370195+sdb9696@users.noreply.github.com>
Date: Fri, 7 Jun 2024 10:52:11 +0100
Subject: [PATCH] Add state features to iot strip sockets (#960)

 Fixes iot strip sockets not creating their own state and on_since features.
---
 kasa/iot/iotstrip.py     | 29 +++++++++++++++++++++++++++++
 kasa/tests/test_strip.py |  8 ++++++++
 2 files changed, 37 insertions(+)

diff --git a/kasa/iot/iotstrip.py b/kasa/iot/iotstrip.py
index 7c6368b0..dde57faa 100755
--- a/kasa/iot/iotstrip.py
+++ b/kasa/iot/iotstrip.py
@@ -10,6 +10,7 @@ from typing import Any
 from ..device_type import DeviceType
 from ..deviceconfig import DeviceConfig
 from ..exceptions import KasaException
+from ..feature import Feature
 from ..module import Module
 from ..protocol import BaseProtocol
 from .iotdevice import (
@@ -125,6 +126,8 @@ class IotStrip(IotDevice):
                 )
                 for child in children
             }
+            for child in self._children.values():
+                await child._initialize_features()
 
         if update_children and self.has_emeter:
             for plug in self.children:
@@ -250,6 +253,32 @@ class IotStripPlug(IotPlug):
         await super()._initialize_modules()
         self.add_module("time", Time(self, "time"))
 
+    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,
+                id="on_since",
+                name="On since",
+                attribute_getter="on_since",
+                icon="mdi:clock",
+            )
+        )
+        # If the strip plug has it's own modules we should call initialize
+        # features for the modules here. However the _initialize_modules function
+        # above does not seem to be called.
+
     async def update(self, update_children: bool = True):
         """Query the device to update the data.
 
diff --git a/kasa/tests/test_strip.py b/kasa/tests/test_strip.py
index e5285acc..4c576d1b 100644
--- a/kasa/tests/test_strip.py
+++ b/kasa/tests/test_strip.py
@@ -88,6 +88,14 @@ async def test_get_plug_by_index(dev: IotStrip):
         dev.get_plug_by_index(len(dev.children))
 
 
+@strip
+async def test_plug_features(dev: IotStrip):
+    """Test the child plugs have default features."""
+    for child in dev.children:
+        assert "state" in child.features
+        assert "on_since" in child.features
+
+
 @pytest.mark.skip("this test will wear out your relays")
 async def test_all_binary_states(dev):
     # test every binary state