From 06fe4e71cf7ec019c0d214c823479519e474c1d3 Mon Sep 17 00:00:00 2001
From: Teemu Rytilahti <tpr@iki.fi>
Date: Fri, 23 Feb 2024 02:32:38 +0100
Subject: [PATCH] Improve smartdevice update module

* Expose current and latest firmware as features
* Implement update loop that blocks until the update is complete
---
 kasa/smart/modules/firmware.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/kasa/smart/modules/firmware.py b/kasa/smart/modules/firmware.py
index 16704ef0..b1302b37 100644
--- a/kasa/smart/modules/firmware.py
+++ b/kasa/smart/modules/firmware.py
@@ -1,6 +1,7 @@
 """Implementation of firmware module."""
 
 from __future__ import annotations
+import asyncio
 
 import asyncio
 import logging
@@ -16,6 +17,11 @@ from ...exceptions import SmartErrorCode
 from ...feature import Feature
 from ..smartmodule import SmartModule
 
+# When support for cpython older than 3.11 is dropped
+# async_timeout can be replaced with asyncio.timeout
+from async_timeout import timeout as asyncio_timeout
+
+
 if TYPE_CHECKING:
     from ..smartdevice import SmartDevice
 
@@ -99,6 +105,12 @@ class Firmware(SmartModule):
                 category=Feature.Category.Info,
             )
         )
+        self._add_feature(
+            Feature(device, "Current firmware version", container=self, attribute_getter="current_firmware")
+        )
+        self._add_feature(
+            Feature(device, "Available firmware version", container=self, attribute_getter="latest_firmware")
+        )
 
     def query(self) -> dict:
         """Query to execute during the update cycle."""
@@ -112,6 +124,7 @@ class Firmware(SmartModule):
         """Return the current firmware version."""
         return self._device.hw_info["sw_ver"]
 
+
     @property
     def latest_firmware(self) -> str:
         """Return the latest firmware version."""