Support for EnergyStorageTrait for vacuum cleaners (#55134)

homeassistant/components/google_assistant/trait.py

@@ -28,6 +28,7 @@ from homeassistant.components.lock import STATE_JAMMED, STATE_UNLOCKING
 from homeassistant.components.media_player.const import MEDIA_TYPE_CHANNEL
 from homeassistant.const import (
     ATTR_ASSUMED_STATE,
+    ATTR_BATTERY_LEVEL,
     ATTR_CODE,
     ATTR_DEVICE_CLASS,
     ATTR_ENTITY_ID,
@@ -106,6 +107,7 @@ TRAIT_TRANSPORT_CONTROL = f"{PREFIX_TRAITS}TransportControl"
 TRAIT_MEDIA_STATE = f"{PREFIX_TRAITS}MediaState"
 TRAIT_CHANNEL = f"{PREFIX_TRAITS}Channel"
 TRAIT_LOCATOR = f"{PREFIX_TRAITS}Locator"
+TRAIT_ENERGYSTORAGE = f"{PREFIX_TRAITS}EnergyStorage"
 
 PREFIX_COMMANDS = "action.devices.commands."
 COMMAND_ONOFF = f"{PREFIX_COMMANDS}OnOff"
@@ -148,6 +150,7 @@ COMMAND_REVERSE = f"{PREFIX_COMMANDS}Reverse"
 COMMAND_SET_HUMIDITY = f"{PREFIX_COMMANDS}SetHumidity"
 COMMAND_SELECT_CHANNEL = f"{PREFIX_COMMANDS}selectChannel"
 COMMAND_LOCATE = f"{PREFIX_COMMANDS}Locate"
+COMMAND_CHARGE = f"{PREFIX_COMMANDS}Charge"
 
 TRAITS = []
 
@@ -609,6 +612,58 @@ class LocatorTrait(_Trait):
         )
 
 
+class EnergyStorageTrait(_Trait):
+    """Trait to offer EnergyStorage functionality.
+
+    https://developers.google.com/actions/smarthome/traits/energystorage
+    """
+
+    name = TRAIT_ENERGYSTORAGE
+    commands = [COMMAND_CHARGE]
+
+    @staticmethod
+    def supported(domain, features, device_class, _):
+        """Test if state is supported."""
+        return domain == vacuum.DOMAIN and features & vacuum.SUPPORT_BATTERY
+
+    def sync_attributes(self):
+        """Return EnergyStorage attributes for a sync request."""
+        return {
+            "isRechargeable": True,
+            "queryOnlyEnergyStorage": True,
+        }
+
+    def query_attributes(self):
+        """Return EnergyStorage query attributes."""
+        battery_level = self.state.attributes.get(ATTR_BATTERY_LEVEL)
+        if battery_level == 100:
+            descriptive_capacity_remaining = "FULL"
+        elif 75 <= battery_level < 100:
+            descriptive_capacity_remaining = "HIGH"
+        elif 50 <= battery_level < 75:
+            descriptive_capacity_remaining = "MEDIUM"
+        elif 25 <= battery_level < 50:
+            descriptive_capacity_remaining = "LOW"
+        elif 0 <= battery_level < 25:
+            descriptive_capacity_remaining = "CRITICALLY_LOW"
+        return {
+            "descriptiveCapacityRemaining": descriptive_capacity_remaining,
+            "capacityRemaining": [{"rawValue": battery_level, "unit": "PERCENTAGE"}],
+            "capacityUntilFull": [
+                {"rawValue": 100 - battery_level, "unit": "PERCENTAGE"}
+            ],
+            "isCharging": self.state.state == vacuum.STATE_DOCKED,
+            "isPluggedIn": self.state.state == vacuum.STATE_DOCKED,
+        }
+
+    async def execute(self, command, data, params, challenge):
+        """Execute a dock command."""
+        raise SmartHomeError(
+            ERR_FUNCTION_NOT_SUPPORTED,
+            "Controlling charging of a vacuum is not yet supported",
+        )
+
+
 @register_trait
 class StartStopTrait(_Trait):
     """Trait to offer StartStop functionality.

tests/components/google_assistant/test_trait.py

@@ -34,6 +34,7 @@ from homeassistant.components.media_player.const import (
 from homeassistant.config import async_process_ha_core_config
 from homeassistant.const import (
     ATTR_ASSUMED_STATE,
+    ATTR_BATTERY_LEVEL,
     ATTR_DEVICE_CLASS,
     ATTR_ENTITY_ID,
     ATTR_MODE,
@@ -387,6 +388,74 @@ async def test_locate_vacuum(hass):
     assert err.value.code == const.ERR_FUNCTION_NOT_SUPPORTED
 
 
+async def test_energystorage_vacuum(hass):
+    """Test EnergyStorage trait support for vacuum domain."""
+    assert helpers.get_google_type(vacuum.DOMAIN, None) is not None
+    assert trait.EnergyStorageTrait.supported(
+        vacuum.DOMAIN, vacuum.SUPPORT_BATTERY, None, None
+    )
+
+    trt = trait.EnergyStorageTrait(
+        hass,
+        State(
+            "vacuum.bla",
+            vacuum.STATE_DOCKED,
+            {
+                ATTR_SUPPORTED_FEATURES: vacuum.SUPPORT_BATTERY,
+                ATTR_BATTERY_LEVEL: 100,
+            },
+        ),
+        BASIC_CONFIG,
+    )
+
+    assert trt.sync_attributes() == {
+        "isRechargeable": True,
+        "queryOnlyEnergyStorage": True,
+    }
+
+    assert trt.query_attributes() == {
+        "descriptiveCapacityRemaining": "FULL",
+        "capacityRemaining": [{"rawValue": 100, "unit": "PERCENTAGE"}],
+        "capacityUntilFull": [{"rawValue": 0, "unit": "PERCENTAGE"}],
+        "isCharging": True,
+        "isPluggedIn": True,
+    }
+
+    trt = trait.EnergyStorageTrait(
+        hass,
+        State(
+            "vacuum.bla",
+            vacuum.STATE_CLEANING,
+            {
+                ATTR_SUPPORTED_FEATURES: vacuum.SUPPORT_BATTERY,
+                ATTR_BATTERY_LEVEL: 20,
+            },
+        ),
+        BASIC_CONFIG,
+    )
+
+    assert trt.sync_attributes() == {
+        "isRechargeable": True,
+        "queryOnlyEnergyStorage": True,
+    }
+
+    assert trt.query_attributes() == {
+        "descriptiveCapacityRemaining": "CRITICALLY_LOW",
+        "capacityRemaining": [{"rawValue": 20, "unit": "PERCENTAGE"}],
+        "capacityUntilFull": [{"rawValue": 80, "unit": "PERCENTAGE"}],
+        "isCharging": False,
+        "isPluggedIn": False,
+    }
+
+    with pytest.raises(helpers.SmartHomeError) as err:
+        await trt.execute(trait.COMMAND_CHARGE, BASIC_DATA, {"charge": True}, {})
+    assert err.value.code == const.ERR_FUNCTION_NOT_SUPPORTED
+
+    with pytest.raises(helpers.SmartHomeError) as err:
+        await trt.execute(trait.COMMAND_CHARGE, BASIC_DATA, {"charge": False}, {})
+    assert err.value.code == const.ERR_FUNCTION_NOT_SUPPORTED
+
+
 async def test_startstop_vacuum(hass):
     """Test startStop trait support for vacuum domain."""
     assert helpers.get_google_type(vacuum.DOMAIN, None) is not None