Added vacuum dock and pause/unpause traits (#17657)

homeassistant/components/google_assistant/trait.py

@@ -13,6 +13,7 @@ from homeassistant.components import (
     scene,
     script,
     switch,
+    vacuum,
 )
 from homeassistant.const import (
     ATTR_ENTITY_ID,
@@ -21,6 +22,7 @@ from homeassistant.const import (
     STATE_OFF,
     TEMP_CELSIUS,
     TEMP_FAHRENHEIT,
+    ATTR_SUPPORTED_FEATURES,
 )
 from homeassistant.util import color as color_util, temperature as temp_util
 
@@ -31,6 +33,8 @@ _LOGGER = logging.getLogger(__name__)
 
 PREFIX_TRAITS = 'action.devices.traits.'
 TRAIT_ONOFF = PREFIX_TRAITS + 'OnOff'
+TRAIT_DOCK = PREFIX_TRAITS + 'Dock'
+TRAIT_STARTSTOP = PREFIX_TRAITS + 'StartStop'
 TRAIT_BRIGHTNESS = PREFIX_TRAITS + 'Brightness'
 TRAIT_COLOR_SPECTRUM = PREFIX_TRAITS + 'ColorSpectrum'
 TRAIT_COLOR_TEMP = PREFIX_TRAITS + 'ColorTemperature'
@@ -39,6 +43,9 @@ TRAIT_TEMPERATURE_SETTING = PREFIX_TRAITS + 'TemperatureSetting'
 
 PREFIX_COMMANDS = 'action.devices.commands.'
 COMMAND_ONOFF = PREFIX_COMMANDS + 'OnOff'
+COMMAND_DOCK = PREFIX_COMMANDS + 'Dock'
+COMMAND_STARTSTOP = PREFIX_COMMANDS + 'StartStop'
+COMMAND_PAUSEUNPAUSE = PREFIX_COMMANDS + 'PauseUnpause'
 COMMAND_BRIGHTNESS_ABSOLUTE = PREFIX_COMMANDS + 'BrightnessAbsolute'
 COMMAND_COLOR_ABSOLUTE = PREFIX_COMMANDS + 'ColorAbsolute'
 COMMAND_ACTIVATE_SCENE = PREFIX_COMMANDS + 'ActivateScene'
@@ -392,6 +399,96 @@ class SceneTrait(_Trait):
             }, blocking=self.state.domain != script.DOMAIN)
 
 
+@register_trait
+class DockTrait(_Trait):
+    """Trait to offer dock functionality.
+
+    https://developers.google.com/actions/smarthome/traits/dock
+    """
+
+    name = TRAIT_DOCK
+    commands = [
+        COMMAND_DOCK
+    ]
+
+    @staticmethod
+    def supported(domain, features):
+        """Test if state is supported."""
+        return domain == vacuum.DOMAIN
+
+    def sync_attributes(self):
+        """Return dock attributes for a sync request."""
+        return {}
+
+    def query_attributes(self):
+        """Return dock query attributes."""
+        return {'isDocked': self.state.state == vacuum.STATE_DOCKED}
+
+    async def execute(self, command, params):
+        """Execute a dock command."""
+        await self.hass.services.async_call(
+            self.state.domain, vacuum.SERVICE_RETURN_TO_BASE, {
+                ATTR_ENTITY_ID: self.state.entity_id
+            }, blocking=True)
+
+
+@register_trait
+class StartStopTrait(_Trait):
+    """Trait to offer StartStop functionality.
+
+    https://developers.google.com/actions/smarthome/traits/startstop
+    """
+
+    name = TRAIT_STARTSTOP
+    commands = [
+        COMMAND_STARTSTOP,
+        COMMAND_PAUSEUNPAUSE
+    ]
+
+    @staticmethod
+    def supported(domain, features):
+        """Test if state is supported."""
+        return domain == vacuum.DOMAIN
+
+    def sync_attributes(self):
+        """Return StartStop attributes for a sync request."""
+        return {'pausable':
+                self.state.attributes.get(ATTR_SUPPORTED_FEATURES, 0)
+                & vacuum.SUPPORT_PAUSE != 0}
+
+    def query_attributes(self):
+        """Return StartStop query attributes."""
+        return {
+            'isRunning': self.state.state == vacuum.STATE_CLEANING,
+            'isPaused': self.state.state == vacuum.STATE_PAUSED,
+        }
+
+    async def execute(self, command, params):
+        """Execute a StartStop command."""
+        if command == COMMAND_STARTSTOP:
+            if params['start']:
+                await self.hass.services.async_call(
+                    self.state.domain, vacuum.SERVICE_START, {
+                        ATTR_ENTITY_ID: self.state.entity_id
+                    }, blocking=True)
+            else:
+                await self.hass.services.async_call(
+                    self.state.domain, vacuum.SERVICE_STOP, {
+                        ATTR_ENTITY_ID: self.state.entity_id
+                    }, blocking=True)
+        elif command == COMMAND_PAUSEUNPAUSE:
+            if params['pause']:
+                await self.hass.services.async_call(
+                    self.state.domain, vacuum.SERVICE_PAUSE, {
+                        ATTR_ENTITY_ID: self.state.entity_id
+                    }, blocking=True)
+            else:
+                await self.hass.services.async_call(
+                    self.state.domain, vacuum.SERVICE_START, {
+                        ATTR_ENTITY_ID: self.state.entity_id
+                    }, blocking=True)
+
+
 @register_trait
 class TemperatureSettingTrait(_Trait):
     """Trait to offer handling both temperature point and modes functionality.