add sensor for battery level (#383)

DeebotUniverse · Jul 16, 2023 · 6a55ca3 · 6a55ca3
1 parent efbd505
commit 6a55ca3
Showing 1 changed file with 16 additions and 1 deletion.
diff --git a/custom_components/deebot/sensor.py b/custom_components/deebot/sensor.py
@@ -5,6 +5,7 @@
 from typing import TypeVar
 
 from deebot_client.events import (
+    BatteryEvent,
     CleanLogEvent,
     ErrorEvent,
     Event,
@@ -16,19 +17,22 @@
 from deebot_client.events.event_bus import EventListener
 from deebot_client.vacuum_bot import VacuumBot
 from homeassistant.components.sensor import (
+    SensorDeviceClass,
     SensorEntity,
     SensorEntityDescription,
     SensorStateClass,
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     AREA_SQUARE_METERS,
+    ATTR_BATTERY_LEVEL,
     CONF_DESCRIPTION,
+    PERCENTAGE,
     TIME_HOURS,
     TIME_MINUTES,
+    EntityCategory,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
@@ -125,6 +129,17 @@ async def async_setup_entry(
                     TotalStatsEvent,
                     lambda e: e.cleanings,
                 ),
+                DeebotGenericSensor(
+                    vacbot,
+                    SensorEntityDescription(
+                        key=ATTR_BATTERY_LEVEL,
+                        native_unit_of_measurement=PERCENTAGE,
+                        device_class=SensorDeviceClass.BATTERY,
+                        entity_category=EntityCategory.DIAGNOSTIC,
+                    ),
+                    BatteryEvent,
+                    lambda b: b.value,
+                ),
             ]
         )