Add sensor for battery

custom_components/robovac/__init__.py

@@ -24,12 +24,12 @@ from .const import CONF_VACS, DOMAIN
 
 from .tuyalocaldiscovery import TuyaLocalDiscovery
 
-PLATFORM = Platform.VACUUM
+PLATFORMS = [Platform.VACUUM, Platform.SENSOR]
 _LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup(hass, entry) -> bool:
-    hass.data.setdefault(DOMAIN, {})
+    hass.data.setdefault(DOMAIN, {CONF_VACS:{}})
 
     async def update_device(device):
         entry = async_get_config_entry_for_device(hass, device["gwId"])
@@ -70,15 +70,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Eufy Robovac from a config entry."""
     entry.async_on_unload(entry.add_update_listener(update_listener))
 
-    await hass.config_entries.async_forward_entry_setup(entry, PLATFORM)
+    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    if unload_ok := await hass.config_entries.async_forward_entry_unload(
+    if unload_ok := await hass.config_entries.async_unload_platforms(
-        entry, PLATFORM
+        entry, PLATFORMS
     ):
         """Nothing"""
     return unload_ok

custom_components/robovac/const.py

@@ -3,3 +3,4 @@
 DOMAIN = "robovac"
 CONF_VACS = "vacuums"
 CONF_AUTODISCOVERY = "autodiscovery"
+REFRESH_RATE = 20

custom_components/robovac/sensor.py

@@ -0,0 +1,63 @@
+import logging
+from datetime import timedelta
+
+from homeassistant.components.sensor import SensorDeviceClass, SensorEntity
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import PERCENTAGE, EntityCategory, CONF_NAME, CONF_ID
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.device_registry import DeviceInfo
+
+from .const import CONF_VACS, DOMAIN, REFRESH_RATE
+
+_LOGGER = logging.getLogger(__name__)
+
+BATTERY = "Battery"
+SCAN_INTERVAL = timedelta(seconds=REFRESH_RATE)
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Initialize my test integration 2 config entry."""
+    vacuums = config_entry.data[CONF_VACS]
+    for item in vacuums:
+        item = vacuums[item]
+        entity = RobovacSensorEntity(item)
+        async_add_entities([entity])
+
+class RobovacSensorEntity(SensorEntity):
+    _attr_has_entity_name = True
+    _attr_device_class = SensorDeviceClass.BATTERY
+    _attr_entity_category = EntityCategory.DIAGNOSTIC
+    _attr_native_unit_of_measurement = PERCENTAGE
+    _attr_available = False
+
+    def __init__(self, item):
+        self.robovac = item
+        self.robovac_id = item[CONF_ID]
+        self._attr_unique_id = item[CONF_ID]
+        self._battery_level = None
+
+        self._attr_device_info = DeviceInfo(
+            identifiers={(DOMAIN, item[CONF_ID])},
+            name=item[CONF_NAME]
+        )
+
+    def update(self):
+        try:
+            self._battery_level = self.hass.data[DOMAIN][CONF_VACS][self.robovac_id].battery_level
+            self._attr_available = True
+        except:
+            _LOGGER.debug("Failed to get battery level for {}".format(self.robovac_id))
+            self._battery_level = None
+            self._attr_available = False
+    
+    @property
+    def native_value(self) -> str | None:
+        """Return the state."""
+        if self._battery_level is not None:
+            return self._battery_level
+        return None
+

custom_components/robovac/vacuum.py

@@ -56,7 +56,7 @@ from homeassistant.const import (
 )
 
 from .tuyalocalapi import TuyaException
-from .const import CONF_VACS, DOMAIN
+from .const import CONF_VACS, DOMAIN, REFRESH_RATE
 
 from .errors import getErrorMessage
 from .robovac import (
@@ -84,7 +84,6 @@ ATTR_CONSUMABLES = "consumables"
 ATTR_MODE = "mode"
 
 _LOGGER = logging.getLogger(__name__)
-REFRESH_RATE = 20
 SCAN_INTERVAL = timedelta(seconds=REFRESH_RATE)
 UPDATE_RETRIES = 3
 
@@ -115,6 +114,7 @@ async def async_setup_entry(
     for item in vacuums:
         item = vacuums[item]
         entity = RoboVacEntity(item)
+        hass.data[DOMAIN][CONF_VACS][item[CONF_ID]] = entity
         async_add_entities([entity])