Added vacuum dock and pause/unpause traits

homeassistant/components/google_assistant/const.py

@@ -15,14 +15,15 @@
 DEFAULT_EXPOSE_BY_DEFAULT = True
 DEFAULT_EXPOSED_DOMAINS = [
     'climate', 'cover', 'fan', 'group', 'input_boolean', 'light',
-    'media_player', 'scene', 'script', 'switch'
+    'media_player', 'scene', 'script', 'switch', 'vacuum',
 ]
 CLIMATE_MODE_HEATCOOL = 'heatcool'
 CLIMATE_SUPPORTED_MODES = {'heat', 'cool', 'off', 'on', CLIMATE_MODE_HEATCOOL}
 
 PREFIX_TYPES = 'action.devices.types.'
 TYPE_LIGHT = PREFIX_TYPES + 'LIGHT'
 TYPE_SWITCH = PREFIX_TYPES + 'SWITCH'
+TYPE_VACUUM = PREFIX_TYPES + 'VACUUM'
 TYPE_SCENE = PREFIX_TYPES + 'SCENE'
 TYPE_THERMOSTAT = PREFIX_TYPES + 'THERMOSTAT'
 

homeassistant/components/google_assistant/smart_home.py

@@ -19,11 +19,13 @@
     scene,
     script,
     switch,
+    vacuum,
 )
 
 from . import trait
 from .const import (
-    TYPE_LIGHT, TYPE_SCENE, TYPE_SWITCH, TYPE_THERMOSTAT,
+    TYPE_LIGHT, TYPE_SCENE, TYPE_SWITCH, TYPE_VACUUM,
+    TYPE_THERMOSTAT,
     CONF_ALIASES, CONF_ROOM_HINT,
     ERR_NOT_SUPPORTED, ERR_PROTOCOL_ERROR, ERR_DEVICE_OFFLINE,
     ERR_UNKNOWN_ERROR
@@ -44,6 +46,7 @@
     scene.DOMAIN: TYPE_SCENE,
     script.DOMAIN: TYPE_SCENE,
     switch.DOMAIN: TYPE_SWITCH,
+    vacuum.DOMAIN: TYPE_VACUUM,
 }
 
 

homeassistant/components/google_assistant/trait.py

@@ -13,6 +13,7 @@
     scene,
     script,
     switch,
+    vacuum,
 )
 from homeassistant.const import (
     ATTR_ENTITY_ID,
@@ -21,6 +22,7 @@
     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 @@
 
 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 @@
 
 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 @@ async def execute(self, command, params):
             }, 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.

tests/components/google_assistant/__init__.py

@@ -184,15 +184,13 @@
         'name': 'Living Room Fan'
     },
     'traits': ['action.devices.traits.OnOff'],
-    'type': 'action.devices.types.SWITCH',
     'willReportState': False
 }, {
     'id': 'fan.ceiling_fan',
     'name': {
         'name': 'Ceiling Fan'
     },
     'traits': ['action.devices.traits.OnOff'],
-    'type': 'action.devices.types.SWITCH',
     'willReportState': False
 }, {
     'id': 'group.all_fans',

tests/components/google_assistant/test_trait.py

@@ -3,7 +3,7 @@
 
 from homeassistant.const import (
     STATE_ON, STATE_OFF, ATTR_ENTITY_ID, SERVICE_TURN_ON, SERVICE_TURN_OFF,
-    TEMP_CELSIUS, TEMP_FAHRENHEIT)
+    TEMP_CELSIUS, TEMP_FAHRENHEIT, ATTR_SUPPORTED_FEATURES)
 from homeassistant.core import State, DOMAIN as HA_DOMAIN
 from homeassistant.components import (
     climate,
@@ -15,6 +15,7 @@
     scene,
     script,
     switch,
+    vacuum,
 )
 from homeassistant.components.google_assistant import trait, helpers, const
 from homeassistant.util import color
@@ -357,6 +358,75 @@ async def test_onoff_media_player(hass):
     }
 
 
+async def test_dock_vacuum(hass):
+    """Test dock trait support for vacuum domain."""
+    assert trait.DockTrait.supported(vacuum.DOMAIN, 0)
+
+    trt = trait.DockTrait(hass, State('vacuum.bla', vacuum.STATE_IDLE))
+
+    assert trt.sync_attributes() == {}
+
+    assert trt.query_attributes() == {
+        'isDocked': False
+    }
+
+    calls = async_mock_service(hass, vacuum.DOMAIN,
+                               vacuum.SERVICE_RETURN_TO_BASE)
+    await trt.execute(trait.COMMAND_DOCK, {})
+    assert len(calls) == 1
+    assert calls[0].data == {
+        ATTR_ENTITY_ID: 'vacuum.bla',
+    }
+
+
+async def test_startstop_vacuum(hass):
+    """Test startStop trait support for vacuum domain."""
+    assert trait.StartStopTrait.supported(vacuum.DOMAIN, 0)
+
+    trt = trait.StartStopTrait(hass, State('vacuum.bla', vacuum.STATE_PAUSED, {
+        ATTR_SUPPORTED_FEATURES: vacuum.SUPPORT_PAUSE,
+    }))
+
+    assert trt.sync_attributes() == {'pausable': True}
+
+    assert trt.query_attributes() == {
+        'isRunning': False,
+        'isPaused': True
+    }
+
+    start_calls = async_mock_service(hass, vacuum.DOMAIN,
+                                     vacuum.SERVICE_START)
+    await trt.execute(trait.COMMAND_STARTSTOP, {'start': True})
+    assert len(start_calls) == 1
+    assert start_calls[0].data == {
+        ATTR_ENTITY_ID: 'vacuum.bla',
+    }
+
+    stop_calls = async_mock_service(hass, vacuum.DOMAIN,
+                                    vacuum.SERVICE_STOP)
+    await trt.execute(trait.COMMAND_STARTSTOP, {'start': False})
+    assert len(stop_calls) == 1
+    assert stop_calls[0].data == {
+        ATTR_ENTITY_ID: 'vacuum.bla',
+    }
+
+    pause_calls = async_mock_service(hass, vacuum.DOMAIN,
+                                     vacuum.SERVICE_PAUSE)
+    await trt.execute(trait.COMMAND_PAUSEUNPAUSE, {'pause': True})
+    assert len(pause_calls) == 1
+    assert pause_calls[0].data == {
+        ATTR_ENTITY_ID: 'vacuum.bla',
+    }
+
+    unpause_calls = async_mock_service(hass, vacuum.DOMAIN,
+                                       vacuum.SERVICE_START)
+    await trt.execute(trait.COMMAND_PAUSEUNPAUSE, {'pause': False})
+    assert len(unpause_calls) == 1
+    assert unpause_calls[0].data == {
+        ATTR_ENTITY_ID: 'vacuum.bla',
+    }
+
+
 async def test_color_spectrum_light(hass):
     """Test ColorSpectrum trait support for light domain."""
     assert not trait.ColorSpectrumTrait.supported(light.DOMAIN, 0)