From 9930311b542c2a8b4d153d880a863956c7c605af Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sat, 7 Oct 2023 09:18:47 -1000
Subject: [PATCH] Parse features only during updates (#527)

Every time emeter functions were called features had to be re-parsed. For power strips, thats a lot of re-parses. Only parse them when we update.
---
 kasa/smartdevice.py            | 32 ++++++++++++++++++++++----------
 kasa/smartstrip.py             |  2 +-
 kasa/tests/test_smartdevice.py |  2 ++
 3 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/kasa/smartdevice.py b/kasa/smartdevice.py
index bdef809a..7463c6db 100755
--- a/kasa/smartdevice.py
+++ b/kasa/smartdevice.py
@@ -92,6 +92,12 @@ def requires_update(f):
     return wrapped
 
 
+@functools.lru_cache
+def _parse_features(features: str) -> Set[str]:
+    """Parse features string."""
+    return set(features.split(":"))
+
+
 class SmartDevice:
     """Base class for all supported device types.
 
@@ -213,6 +219,7 @@ class SmartDevice:
         #       are not accessed incorrectly.
         self._last_update: Any = None
         self._sys_info: Any = None  # TODO: this is here to avoid changing tests
+        self._features: Set[str] = set()
         self.modules: Dict[str, Any] = {}
 
         self.children: List["SmartDevice"] = []
@@ -284,11 +291,7 @@ class SmartDevice:
     @requires_update
     def features(self) -> Set[str]:
         """Return a set of features that the device supports."""
-        try:
-            return set(self.sys_info["feature"].split(":"))
-        except KeyError:
-            _LOGGER.debug("Device does not have feature information")
-            return set()
+        return self._features
 
     @property  # type: ignore
     @requires_update
@@ -321,11 +324,12 @@ class SmartDevice:
         # See #105, #120, #161
         if self._last_update is None:
             _LOGGER.debug("Performing the initial update to obtain sysinfo")
-            self._last_update = await self.protocol.query(req)
-            self._sys_info = self._last_update["system"]["get_sysinfo"]
+            response = await self.protocol.query(req)
+            self._last_update = response
+            self._set_sys_info(response["system"]["get_sysinfo"])
 
         await self._modular_update(req)
-        self._sys_info = self._last_update["system"]["get_sysinfo"]
+        self._set_sys_info(self._last_update["system"]["get_sysinfo"])
 
     async def _modular_update(self, req: dict) -> None:
         """Execute an update query."""
@@ -366,10 +370,18 @@ class SmartDevice:
             update = {**update, **response}
         self._last_update = update
 
-    def update_from_discover_info(self, info):
+    def update_from_discover_info(self, info: Dict[str, Any]) -> None:
         """Update state from info from the discover call."""
         self._last_update = info
-        self._sys_info = info["system"]["get_sysinfo"]
+        self._set_sys_info(info["system"]["get_sysinfo"])
+
+    def _set_sys_info(self, sys_info: Dict[str, Any]) -> None:
+        """Set sys_info."""
+        self._sys_info = sys_info
+        if features := sys_info.get("feature"):
+            self._features = _parse_features(features)
+        else:
+            self._features = set()
 
     @property  # type: ignore
     @requires_update
diff --git a/kasa/smartstrip.py b/kasa/smartstrip.py
index 2a55b2a8..2e5d80a4 100755
--- a/kasa/smartstrip.py
+++ b/kasa/smartstrip.py
@@ -259,7 +259,7 @@ class SmartStripPlug(SmartPlug):
         self.parent = parent
         self.child_id = child_id
         self._last_update = parent._last_update
-        self._sys_info = parent._sys_info
+        self._set_sys_info(parent.sys_info)
         self._device_type = DeviceType.StripSocket
         self.modules = {}
         self.protocol = parent.protocol  # Must use the same connection as the parent
diff --git a/kasa/tests/test_smartdevice.py b/kasa/tests/test_smartdevice.py
index ec4e3d56..aeb1fe03 100644
--- a/kasa/tests/test_smartdevice.py
+++ b/kasa/tests/test_smartdevice.py
@@ -37,6 +37,7 @@ async def test_invalid_connection(dev):
 async def test_initial_update_emeter(dev, mocker):
     """Test that the initial update performs second query if emeter is available."""
     dev._last_update = None
+    dev._features = set()
     spy = mocker.spy(dev.protocol, "query")
     await dev.update()
     # Devices with small buffers may require 3 queries
@@ -48,6 +49,7 @@ async def test_initial_update_emeter(dev, mocker):
 async def test_initial_update_no_emeter(dev, mocker):
     """Test that the initial update performs second query if emeter is available."""
     dev._last_update = None
+    dev._features = set()
     spy = mocker.spy(dev.protocol, "query")
     await dev.update()
     # 2 calls are necessary as some devices crash on unexpected modules