Skip to content

Commit

Permalink
Update entity feature constants google_assistant (home-assistant#98335)
Browse files Browse the repository at this point in the history 
* Update entity feature constants google_assistant

* Update tests

* Direct import

* Some missed constants

* Add fan and cover feature imports
  • Loading branch information
@jbouwh
jbouwh authored Aug 13, 2023
1 parent fa6ffd9 commit e5f7d83
Show file tree
Hide file tree
Showing 3 changed files with 158 additions and 105 deletions.
91 changes: 51 additions & 40 deletions homeassistant/components/google_assistant/trait.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,16 @@
switch,
vacuum,
)
from homeassistant.components.alarm_control_panel import AlarmControlPanelEntityFeature
from homeassistant.components.camera import CameraEntityFeature
from homeassistant.components.climate import ClimateEntityFeature
from homeassistant.components.cover import CoverEntityFeature
from homeassistant.components.fan import FanEntityFeature
from homeassistant.components.humidifier import HumidifierEntityFeature
from homeassistant.components.light import LightEntityFeature
from homeassistant.components.lock import STATE_JAMMED, STATE_UNLOCKING
from homeassistant.components.media_player import MediaType
from homeassistant.components.media_player import MediaPlayerEntityFeature, MediaType
from homeassistant.components.vacuum import VacuumEntityFeature
from homeassistant.const import (
ATTR_ASSUMED_STATE,
ATTR_BATTERY_LEVEL,
Expand Down Expand Up @@ -302,7 +310,7 @@ class CameraStreamTrait(_Trait):
def supported(domain, features, device_class, _):
"""Test if state is supported."""
if domain == camera.DOMAIN:
return features & camera.SUPPORT_STREAM
return features & CameraEntityFeature.STREAM

return False

Expand Down Expand Up @@ -612,7 +620,7 @@ class LocatorTrait(_Trait):
@staticmethod
def supported(domain, features, device_class, _):
"""Test if state is supported."""
return domain == vacuum.DOMAIN and features & vacuum.SUPPORT_LOCATE
return domain == vacuum.DOMAIN and features & VacuumEntityFeature.LOCATE

def sync_attributes(self):
"""Return locator attributes for a sync request."""
Expand Down Expand Up @@ -652,7 +660,7 @@ class EnergyStorageTrait(_Trait):
@staticmethod
def supported(domain, features, device_class, _):
"""Test if state is supported."""
return domain == vacuum.DOMAIN and features & vacuum.SUPPORT_BATTERY
return domain == vacuum.DOMAIN and features & VacuumEntityFeature.BATTERY

def sync_attributes(self):
"""Return EnergyStorage attributes for a sync request."""
Expand Down Expand Up @@ -710,7 +718,7 @@ def supported(domain, features, device_class, _):
if domain == vacuum.DOMAIN:
return True

if domain == cover.DOMAIN and features & cover.SUPPORT_STOP:
if domain == cover.DOMAIN and features & CoverEntityFeature.STOP:
return True

return False
Expand All @@ -721,7 +729,7 @@ def sync_attributes(self):
if domain == vacuum.DOMAIN:
return {
"pausable": self.state.attributes.get(ATTR_SUPPORTED_FEATURES, 0)
& vacuum.SUPPORT_PAUSE
& VacuumEntityFeature.PAUSE
!= 0
}
if domain == cover.DOMAIN:
Expand Down Expand Up @@ -991,7 +999,7 @@ def query_attributes(self):
response["thermostatHumidityAmbient"] = current_humidity

if operation in (climate.HVACMode.AUTO, climate.HVACMode.HEAT_COOL):
if supported & climate.SUPPORT_TARGET_TEMPERATURE_RANGE:
if supported & ClimateEntityFeature.TARGET_TEMPERATURE_RANGE:
response["thermostatTemperatureSetpointHigh"] = round(
TemperatureConverter.convert(
attrs[climate.ATTR_TARGET_TEMP_HIGH],
Expand Down Expand Up @@ -1093,7 +1101,7 @@ async def execute(self, command, data, params, challenge):
supported = self.state.attributes.get(ATTR_SUPPORTED_FEATURES)
svc_data = {ATTR_ENTITY_ID: self.state.entity_id}

if supported & climate.SUPPORT_TARGET_TEMPERATURE_RANGE:
if supported & ClimateEntityFeature.TARGET_TEMPERATURE_RANGE:
svc_data[climate.ATTR_TARGET_TEMP_HIGH] = temp_high
svc_data[climate.ATTR_TARGET_TEMP_LOW] = temp_low
else:
Expand Down Expand Up @@ -1311,11 +1319,11 @@ class ArmDisArmTrait(_Trait):
}

state_to_support = {
STATE_ALARM_ARMED_HOME: alarm_control_panel.const.SUPPORT_ALARM_ARM_HOME,
STATE_ALARM_ARMED_AWAY: alarm_control_panel.const.SUPPORT_ALARM_ARM_AWAY,
STATE_ALARM_ARMED_NIGHT: alarm_control_panel.const.SUPPORT_ALARM_ARM_NIGHT,
STATE_ALARM_ARMED_CUSTOM_BYPASS: alarm_control_panel.const.SUPPORT_ALARM_ARM_CUSTOM_BYPASS,
STATE_ALARM_TRIGGERED: alarm_control_panel.const.SUPPORT_ALARM_TRIGGER,
STATE_ALARM_ARMED_HOME: AlarmControlPanelEntityFeature.ARM_HOME,
STATE_ALARM_ARMED_AWAY: AlarmControlPanelEntityFeature.ARM_AWAY,
STATE_ALARM_ARMED_NIGHT: AlarmControlPanelEntityFeature.ARM_NIGHT,
STATE_ALARM_ARMED_CUSTOM_BYPASS: AlarmControlPanelEntityFeature.ARM_CUSTOM_BYPASS,
STATE_ALARM_TRIGGERED: AlarmControlPanelEntityFeature.TRIGGER,
}

@staticmethod
Expand Down Expand Up @@ -1454,9 +1462,9 @@ def __init__(self, hass, state, config):
def supported(domain, features, device_class, _):
"""Test if state is supported."""
if domain == fan.DOMAIN:
return features & fan.SUPPORT_SET_SPEED
return features & FanEntityFeature.SET_SPEED
if domain == climate.DOMAIN:
return features & climate.SUPPORT_FAN_MODE
return features & ClimateEntityFeature.FAN_MODE
return False

def sync_attributes(self):
Expand All @@ -1468,7 +1476,7 @@ def sync_attributes(self):
if domain == fan.DOMAIN:
reversible = bool(
self.state.attributes.get(ATTR_SUPPORTED_FEATURES, 0)
& fan.SUPPORT_DIRECTION
& FanEntityFeature.DIRECTION
)

result.update(
Expand Down Expand Up @@ -1604,7 +1612,7 @@ class ModesTrait(_Trait):
@staticmethod
def supported(domain, features, device_class, _):
"""Test if state is supported."""
if domain == fan.DOMAIN and features & fan.SUPPORT_PRESET_MODE:
if domain == fan.DOMAIN and features & FanEntityFeature.PRESET_MODE:
return True

if domain == input_select.DOMAIN:
Expand All @@ -1613,16 +1621,16 @@ def supported(domain, features, device_class, _):
if domain == select.DOMAIN:
return True

if domain == humidifier.DOMAIN and features & humidifier.SUPPORT_MODES:
if domain == humidifier.DOMAIN and features & HumidifierEntityFeature.MODES:
return True

if domain == light.DOMAIN and features & light.SUPPORT_EFFECT:
if domain == light.DOMAIN and features & LightEntityFeature.EFFECT:
return True

if domain != media_player.DOMAIN:
return False

return features & media_player.SUPPORT_SELECT_SOUND_MODE
return features & MediaPlayerEntityFeature.SELECT_SOUND_MODE

def _generate(self, name, settings):
"""Generate a list of modes."""
Expand Down Expand Up @@ -1812,7 +1820,7 @@ class InputSelectorTrait(_Trait):
def supported(domain, features, device_class, _):
"""Test if state is supported."""
if domain == media_player.DOMAIN and (
features & media_player.SUPPORT_SELECT_SOURCE
features & MediaPlayerEntityFeature.SELECT_SOURCE
):
return True

Expand Down Expand Up @@ -1910,13 +1918,13 @@ def sync_attributes(self):
response["discreteOnlyOpenClose"] = True
elif (
self.state.domain == cover.DOMAIN
and features & cover.SUPPORT_SET_POSITION == 0
and features & CoverEntityFeature.SET_POSITION == 0
):
response["discreteOnlyOpenClose"] = True

if (
features & cover.SUPPORT_OPEN == 0
and features & cover.SUPPORT_CLOSE == 0
features & CoverEntityFeature.OPEN == 0
and features & CoverEntityFeature.CLOSE == 0
):
response["queryOnlyOpenClose"] = True

Expand Down Expand Up @@ -1985,7 +1993,7 @@ async def execute(self, command, data, params, challenge):
elif position == 100:
service = cover.SERVICE_OPEN_COVER
should_verify = True
elif features & cover.SUPPORT_SET_POSITION:
elif features & CoverEntityFeature.SET_POSITION:
service = cover.SERVICE_SET_COVER_POSITION
if position > 0:
should_verify = True
Expand Down Expand Up @@ -2026,7 +2034,8 @@ def supported(domain, features, device_class, _):
"""Test if trait is supported."""
if domain == media_player.DOMAIN:
return features & (
media_player.SUPPORT_VOLUME_SET | media_player.SUPPORT_VOLUME_STEP
MediaPlayerEntityFeature.VOLUME_SET
| MediaPlayerEntityFeature.VOLUME_STEP
)

return False
Expand All @@ -2035,7 +2044,9 @@ def sync_attributes(self):
"""Return volume attributes for a sync request."""
features = self.state.attributes.get(ATTR_SUPPORTED_FEATURES, 0)
return {
"volumeCanMuteAndUnmute": bool(features & media_player.SUPPORT_VOLUME_MUTE),
"volumeCanMuteAndUnmute": bool(
features & MediaPlayerEntityFeature.VOLUME_MUTE
),
"commandOnlyVolume": self.state.attributes.get(ATTR_ASSUMED_STATE, False),
# Volume amounts in SET_VOLUME and VOLUME_RELATIVE are on a scale
# from 0 to this value.
Expand Down Expand Up @@ -2078,7 +2089,7 @@ async def _execute_set_volume(self, data, params):

if not (
self.state.attributes.get(ATTR_SUPPORTED_FEATURES, 0)
& media_player.SUPPORT_VOLUME_SET
& MediaPlayerEntityFeature.VOLUME_SET
):
raise SmartHomeError(ERR_NOT_SUPPORTED, "Command not supported")

Expand All @@ -2088,13 +2099,13 @@ async def _execute_volume_relative(self, data, params):
relative = params["relativeSteps"]
features = self.state.attributes.get(ATTR_SUPPORTED_FEATURES, 0)

if features & media_player.SUPPORT_VOLUME_SET:
if features & MediaPlayerEntityFeature.VOLUME_SET:
current = self.state.attributes.get(media_player.ATTR_MEDIA_VOLUME_LEVEL)
target = max(0.0, min(1.0, current + relative / 100))

await self._set_volume_absolute(data, target)

elif features & media_player.SUPPORT_VOLUME_STEP:
elif features & MediaPlayerEntityFeature.VOLUME_STEP:
svc = media_player.SERVICE_VOLUME_UP
if relative < 0:
svc = media_player.SERVICE_VOLUME_DOWN
Expand All @@ -2116,7 +2127,7 @@ async def _execute_mute(self, data, params):

if not (
self.state.attributes.get(ATTR_SUPPORTED_FEATURES, 0)
& media_player.SUPPORT_VOLUME_MUTE
& MediaPlayerEntityFeature.VOLUME_MUTE
):
raise SmartHomeError(ERR_NOT_SUPPORTED, "Command not supported")

Expand Down Expand Up @@ -2158,14 +2169,14 @@ def _verify_pin_challenge(data, state, challenge):


MEDIA_COMMAND_SUPPORT_MAPPING = {
COMMAND_MEDIA_NEXT: media_player.SUPPORT_NEXT_TRACK,
COMMAND_MEDIA_PAUSE: media_player.SUPPORT_PAUSE,
COMMAND_MEDIA_PREVIOUS: media_player.SUPPORT_PREVIOUS_TRACK,
COMMAND_MEDIA_RESUME: media_player.SUPPORT_PLAY,
COMMAND_MEDIA_SEEK_RELATIVE: media_player.SUPPORT_SEEK,
COMMAND_MEDIA_SEEK_TO_POSITION: media_player.SUPPORT_SEEK,
COMMAND_MEDIA_SHUFFLE: media_player.SUPPORT_SHUFFLE_SET,
COMMAND_MEDIA_STOP: media_player.SUPPORT_STOP,
COMMAND_MEDIA_NEXT: MediaPlayerEntityFeature.NEXT_TRACK,
COMMAND_MEDIA_PAUSE: MediaPlayerEntityFeature.PAUSE,
COMMAND_MEDIA_PREVIOUS: MediaPlayerEntityFeature.PREVIOUS_TRACK,
COMMAND_MEDIA_RESUME: MediaPlayerEntityFeature.PLAY,
COMMAND_MEDIA_SEEK_RELATIVE: MediaPlayerEntityFeature.SEEK,
COMMAND_MEDIA_SEEK_TO_POSITION: MediaPlayerEntityFeature.SEEK,
COMMAND_MEDIA_SHUFFLE: MediaPlayerEntityFeature.SHUFFLE_SET,
COMMAND_MEDIA_STOP: MediaPlayerEntityFeature.STOP,
}

MEDIA_COMMAND_ATTRIBUTES = {
Expand Down Expand Up @@ -2350,7 +2361,7 @@ def supported(domain, features, device_class, _):
"""Test if state is supported."""
if (
domain == media_player.DOMAIN
and (features & media_player.SUPPORT_PLAY_MEDIA)
and (features & MediaPlayerEntityFeature.PLAY_MEDIA)
and device_class == media_player.MediaPlayerDeviceClass.TV
):
return True
Expand Down
6 changes: 4 additions & 2 deletions tests/components/google_assistant/test_smart_home.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import pytest
from pytest_unordered import unordered

from homeassistant.components import camera
from homeassistant.components.camera import CameraEntityFeature
from homeassistant.components.climate import ATTR_MAX_TEMP, ATTR_MIN_TEMP, HVACMode
from homeassistant.components.demo.binary_sensor import DemoBinarySensor
from homeassistant.components.demo.cover import DemoCover
Expand Down Expand Up @@ -1186,7 +1186,9 @@ async def test_trait_execute_adding_query_data(hass: HomeAssistant) -> None:
{"external_url": "https://example.com"},
)
hass.states.async_set(
"camera.office", "idle", {"supported_features": camera.SUPPORT_STREAM}
"camera.office",
"idle",
{"supported_features": CameraEntityFeature.STREAM},
)

with patch(
Expand Down
Loading

0 comments on commit e5f7d83

Please sign in to comment.