From 8ff6172c30957006ccab163aba635bca7138688e Mon Sep 17 00:00:00 2001
From: Teemu Rytilahti <tpr@iki.fi>
Date: Sat, 30 Nov 2024 17:26:56 +0100
Subject: [PATCH] Add battery level

---
 kasa/smart/modules/vacuum.py | 29 ++++++++++++++++++++++++++---
 1 file changed, 26 insertions(+), 3 deletions(-)

diff --git a/kasa/smart/modules/vacuum.py b/kasa/smart/modules/vacuum.py
index b04fce0d..8d2d240f 100644
--- a/kasa/smart/modules/vacuum.py
+++ b/kasa/smart/modules/vacuum.py
@@ -80,10 +80,23 @@ class Vacuum(SmartModule):
                 type=Feature.Type.Sensor,
             )
         )
+        self._add_feature(
+            Feature(
+                self._device,
+                "battery_level",
+                "Battery level",
+                container=self,
+                attribute_getter="battery",
+                icon="mdi:battery",
+                unit_getter=lambda: "%",
+                category=Feature.Category.Info,
+                type=Feature.Type.Sensor,
+            )
+        )
 
     def query(self) -> dict:
         """Query to execute during the update cycle."""
-        return {"getVacStatus": None}
+        return {"getVacStatus": None, "getBatteryInfo": None}
 
     async def start(self) -> dict:
         """Start cleaning."""
@@ -123,13 +136,23 @@ class Vacuum(SmartModule):
         """Return home / pause returning."""
         return await self.call("setSwitchCharge", {"switch_charge": enabled})
 
+    @property
+    def battery(self) -> int:
+        """Return battery level."""
+        return self.data["getBatteryInfo"]["battery_percentage"]
+
+    @property
+    def _vac_status(self) -> dict:
+        """Return vac status container."""
+        return self.data["getVacStatus"]
+
     @property
     def status(self) -> Status:
         """Return current status."""
-        if self.data.get("err_status"):
+        if self._vac_status.get("err_status"):
             return Status.Error
 
-        status_code = self.data["status"]
+        status_code = self._vac_status["status"]
         try:
             return Status(status_code)
         except ValueError: