diff --git a/homeassistant/components/google_assistant/trait.py b/homeassistant/components/google_assistant/trait.py index 36660820efb..425a394b522 100644 --- a/homeassistant/components/google_assistant/trait.py +++ b/homeassistant/components/google_assistant/trait.py @@ -28,8 +28,16 @@ from homeassistant.components import ( 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, @@ -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 @@ -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.""" @@ -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.""" @@ -710,7 +718,7 @@ class StartStopTrait(_Trait): 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 @@ -721,7 +729,7 @@ class StartStopTrait(_Trait): if domain == vacuum.DOMAIN: return { "pausable": self.state.attributes.get(ATTR_SUPPORTED_FEATURES, 0) - & vacuum.SUPPORT_PAUSE + & VacuumEntityFeature.PAUSE != 0 } if domain == cover.DOMAIN: @@ -991,7 +999,7 @@ class TemperatureSettingTrait(_Trait): 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], @@ -1093,7 +1101,7 @@ class TemperatureSettingTrait(_Trait): 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: @@ -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 @@ -1454,9 +1462,9 @@ class FanSpeedTrait(_Trait): 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): @@ -1468,7 +1476,7 @@ class FanSpeedTrait(_Trait): if domain == fan.DOMAIN: reversible = bool( self.state.attributes.get(ATTR_SUPPORTED_FEATURES, 0) - & fan.SUPPORT_DIRECTION + & FanEntityFeature.DIRECTION ) result.update( @@ -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: @@ -1613,16 +1621,16 @@ class ModesTrait(_Trait): 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.""" @@ -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 @@ -1910,13 +1918,13 @@ class OpenCloseTrait(_Trait): 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 @@ -1985,7 +1993,7 @@ class OpenCloseTrait(_Trait): 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 @@ -2026,7 +2034,8 @@ class VolumeTrait(_Trait): """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 @@ -2035,7 +2044,9 @@ class VolumeTrait(_Trait): """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. @@ -2078,7 +2089,7 @@ class VolumeTrait(_Trait): 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") @@ -2088,13 +2099,13 @@ class VolumeTrait(_Trait): 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 @@ -2116,7 +2127,7 @@ class VolumeTrait(_Trait): 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") @@ -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 = { @@ -2350,7 +2361,7 @@ class ChannelTrait(_Trait): """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 diff --git a/tests/components/google_assistant/test_smart_home.py b/tests/components/google_assistant/test_smart_home.py index 6cfa7965074..bf48564c251 100644 --- a/tests/components/google_assistant/test_smart_home.py +++ b/tests/components/google_assistant/test_smart_home.py @@ -6,7 +6,7 @@ from unittest.mock import ANY, patch 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 @@ -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( diff --git a/tests/components/google_assistant/test_trait.py b/tests/components/google_assistant/test_trait.py index fd6b3a6790b..fcbf16c21c7 100644 --- a/tests/components/google_assistant/test_trait.py +++ b/tests/components/google_assistant/test_trait.py @@ -27,9 +27,22 @@ from homeassistant.components import ( 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.google_assistant import const, error, helpers, trait from homeassistant.components.google_assistant.error import SmartHomeError -from homeassistant.components.media_player import SERVICE_PLAY_MEDIA, MediaType +from homeassistant.components.humidifier import HumidifierEntityFeature +from homeassistant.components.light import LightEntityFeature +from homeassistant.components.lock import LockEntityFeature +from homeassistant.components.media_player import ( + SERVICE_PLAY_MEDIA, + MediaPlayerEntityFeature, + MediaType, +) +from homeassistant.components.vacuum import VacuumEntityFeature from homeassistant.config import async_process_ha_core_config from homeassistant.const import ( ATTR_ASSUMED_STATE, @@ -126,7 +139,7 @@ async def test_camera_stream(hass: HomeAssistant) -> None: ) assert helpers.get_google_type(camera.DOMAIN, None) is not None assert trait.CameraStreamTrait.supported( - camera.DOMAIN, camera.SUPPORT_STREAM, None, None + camera.DOMAIN, CameraEntityFeature.STREAM, None, None ) trt = trait.CameraStreamTrait( @@ -364,7 +377,7 @@ async def test_locate_vacuum(hass: HomeAssistant) -> None: """Test locate trait support for vacuum domain.""" assert helpers.get_google_type(vacuum.DOMAIN, None) is not None assert trait.LocatorTrait.supported( - vacuum.DOMAIN, vacuum.SUPPORT_LOCATE, None, None + vacuum.DOMAIN, VacuumEntityFeature.LOCATE, None, None ) trt = trait.LocatorTrait( @@ -372,7 +385,7 @@ async def test_locate_vacuum(hass: HomeAssistant) -> None: State( "vacuum.bla", vacuum.STATE_IDLE, - {ATTR_SUPPORTED_FEATURES: vacuum.SUPPORT_LOCATE}, + {ATTR_SUPPORTED_FEATURES: VacuumEntityFeature.LOCATE}, ), BASIC_CONFIG, ) @@ -395,7 +408,7 @@ async def test_energystorage_vacuum(hass: HomeAssistant) -> None: """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 + vacuum.DOMAIN, VacuumEntityFeature.BATTERY, None, None ) trt = trait.EnergyStorageTrait( @@ -404,7 +417,7 @@ async def test_energystorage_vacuum(hass: HomeAssistant) -> None: "vacuum.bla", vacuum.STATE_DOCKED, { - ATTR_SUPPORTED_FEATURES: vacuum.SUPPORT_BATTERY, + ATTR_SUPPORTED_FEATURES: VacuumEntityFeature.BATTERY, ATTR_BATTERY_LEVEL: 100, }, ), @@ -430,7 +443,7 @@ async def test_energystorage_vacuum(hass: HomeAssistant) -> None: "vacuum.bla", vacuum.STATE_CLEANING, { - ATTR_SUPPORTED_FEATURES: vacuum.SUPPORT_BATTERY, + ATTR_SUPPORTED_FEATURES: VacuumEntityFeature.BATTERY, ATTR_BATTERY_LEVEL: 20, }, ), @@ -469,7 +482,7 @@ async def test_startstop_vacuum(hass: HomeAssistant) -> None: State( "vacuum.bla", vacuum.STATE_PAUSED, - {ATTR_SUPPORTED_FEATURES: vacuum.SUPPORT_PAUSE}, + {ATTR_SUPPORTED_FEATURES: VacuumEntityFeature.PAUSE}, ), BASIC_CONFIG, ) @@ -502,12 +515,14 @@ async def test_startstop_vacuum(hass: HomeAssistant) -> None: async def test_startstop_cover(hass: HomeAssistant) -> None: """Test startStop trait support for cover domain.""" assert helpers.get_google_type(cover.DOMAIN, None) is not None - assert trait.StartStopTrait.supported(cover.DOMAIN, cover.SUPPORT_STOP, None, None) + assert trait.StartStopTrait.supported( + cover.DOMAIN, CoverEntityFeature.STOP, None, None + ) state = State( "cover.bla", cover.STATE_CLOSED, - {ATTR_SUPPORTED_FEATURES: cover.SUPPORT_STOP}, + {ATTR_SUPPORTED_FEATURES: CoverEntityFeature.STOP}, ) trt = trait.StartStopTrait( @@ -551,7 +566,10 @@ async def test_startstop_cover_assumed(hass: HomeAssistant) -> None: State( "cover.bla", cover.STATE_CLOSED, - {ATTR_SUPPORTED_FEATURES: cover.SUPPORT_STOP, ATTR_ASSUMED_STATE: True}, + { + ATTR_SUPPORTED_FEATURES: CoverEntityFeature.STOP, + ATTR_ASSUMED_STATE: True, + }, ), BASIC_CONFIG, ) @@ -707,7 +725,9 @@ async def test_color_light_temperature_light_bad_temp(hass: HomeAssistant) -> No async def test_light_modes(hass: HomeAssistant) -> None: """Test Light Mode trait.""" assert helpers.get_google_type(light.DOMAIN, None) is not None - assert trait.ModesTrait.supported(light.DOMAIN, light.SUPPORT_EFFECT, None, None) + assert trait.ModesTrait.supported( + light.DOMAIN, LightEntityFeature.EFFECT, None, None + ) trt = trait.ModesTrait( hass, @@ -847,7 +867,7 @@ async def test_temperature_setting_climate_onoff(hass: HomeAssistant) -> None: "climate.bla", climate.HVACMode.AUTO, { - ATTR_SUPPORTED_FEATURES: climate.SUPPORT_TARGET_TEMPERATURE_RANGE, + ATTR_SUPPORTED_FEATURES: ClimateEntityFeature.TARGET_TEMPERATURE_RANGE, climate.ATTR_HVAC_MODES: [ climate.HVACMode.OFF, climate.HVACMode.COOL, @@ -928,7 +948,7 @@ async def test_temperature_setting_climate_range(hass: HomeAssistant) -> None: { climate.ATTR_CURRENT_TEMPERATURE: 70, climate.ATTR_CURRENT_HUMIDITY: 25, - ATTR_SUPPORTED_FEATURES: climate.SUPPORT_TARGET_TEMPERATURE_RANGE, + ATTR_SUPPORTED_FEATURES: ClimateEntityFeature.TARGET_TEMPERATURE_RANGE, climate.ATTR_HVAC_MODES: [ STATE_OFF, climate.HVACMode.COOL, @@ -1040,7 +1060,7 @@ async def test_temperature_setting_climate_setpoint(hass: HomeAssistant) -> None "climate.bla", climate.HVACMode.COOL, { - ATTR_SUPPORTED_FEATURES: climate.SUPPORT_TARGET_TEMPERATURE, + ATTR_SUPPORTED_FEATURES: ClimateEntityFeature.TARGET_TEMPERATURE, climate.ATTR_HVAC_MODES: [STATE_OFF, climate.HVACMode.COOL], climate.ATTR_MIN_TEMP: 10, climate.ATTR_MAX_TEMP: 30, @@ -1230,8 +1250,10 @@ async def test_humidity_setting_humidifier_setpoint(hass: HomeAssistant) -> None async def test_lock_unlock_lock(hass: HomeAssistant) -> None: """Test LockUnlock trait locking support for lock domain.""" assert helpers.get_google_type(lock.DOMAIN, None) is not None - assert trait.LockUnlockTrait.supported(lock.DOMAIN, lock.SUPPORT_OPEN, None, None) - assert trait.LockUnlockTrait.might_2fa(lock.DOMAIN, lock.SUPPORT_OPEN, None) + assert trait.LockUnlockTrait.supported( + lock.DOMAIN, LockEntityFeature.OPEN, None, None + ) + assert trait.LockUnlockTrait.might_2fa(lock.DOMAIN, LockEntityFeature.OPEN, None) trt = trait.LockUnlockTrait( hass, State("lock.front_door", lock.STATE_LOCKED), PIN_CONFIG @@ -1254,8 +1276,10 @@ async def test_lock_unlock_lock(hass: HomeAssistant) -> None: async def test_lock_unlock_unlocking(hass: HomeAssistant) -> None: """Test LockUnlock trait locking support for lock domain.""" assert helpers.get_google_type(lock.DOMAIN, None) is not None - assert trait.LockUnlockTrait.supported(lock.DOMAIN, lock.SUPPORT_OPEN, None, None) - assert trait.LockUnlockTrait.might_2fa(lock.DOMAIN, lock.SUPPORT_OPEN, None) + assert trait.LockUnlockTrait.supported( + lock.DOMAIN, LockEntityFeature.OPEN, None, None + ) + assert trait.LockUnlockTrait.might_2fa(lock.DOMAIN, LockEntityFeature.OPEN, None) trt = trait.LockUnlockTrait( hass, State("lock.front_door", lock.STATE_UNLOCKING), PIN_CONFIG @@ -1269,8 +1293,10 @@ async def test_lock_unlock_unlocking(hass: HomeAssistant) -> None: async def test_lock_unlock_lock_jammed(hass: HomeAssistant) -> None: """Test LockUnlock trait locking support for lock domain that jams.""" assert helpers.get_google_type(lock.DOMAIN, None) is not None - assert trait.LockUnlockTrait.supported(lock.DOMAIN, lock.SUPPORT_OPEN, None, None) - assert trait.LockUnlockTrait.might_2fa(lock.DOMAIN, lock.SUPPORT_OPEN, None) + assert trait.LockUnlockTrait.supported( + lock.DOMAIN, LockEntityFeature.OPEN, None, None + ) + assert trait.LockUnlockTrait.might_2fa(lock.DOMAIN, LockEntityFeature.OPEN, None) trt = trait.LockUnlockTrait( hass, State("lock.front_door", lock.STATE_JAMMED), PIN_CONFIG @@ -1293,7 +1319,9 @@ async def test_lock_unlock_lock_jammed(hass: HomeAssistant) -> None: async def test_lock_unlock_unlock(hass: HomeAssistant) -> None: """Test LockUnlock trait unlocking support for lock domain.""" assert helpers.get_google_type(lock.DOMAIN, None) is not None - assert trait.LockUnlockTrait.supported(lock.DOMAIN, lock.SUPPORT_OPEN, None, None) + assert trait.LockUnlockTrait.supported( + lock.DOMAIN, LockEntityFeature.OPEN, None, None + ) trt = trait.LockUnlockTrait( hass, State("lock.front_door", lock.STATE_LOCKED), PIN_CONFIG @@ -1363,8 +1391,8 @@ async def test_arm_disarm_arm_away(hass: HomeAssistant) -> None: STATE_ALARM_ARMED_AWAY, { alarm_control_panel.ATTR_CODE_ARM_REQUIRED: True, - ATTR_SUPPORTED_FEATURES: alarm_control_panel.const.SUPPORT_ALARM_ARM_HOME - | alarm_control_panel.const.SUPPORT_ALARM_ARM_AWAY, + ATTR_SUPPORTED_FEATURES: AlarmControlPanelEntityFeature.ARM_HOME + | AlarmControlPanelEntityFeature.ARM_AWAY, }, ), PIN_CONFIG, @@ -1526,8 +1554,8 @@ async def test_arm_disarm_disarm(hass: HomeAssistant) -> None: STATE_ALARM_DISARMED, { alarm_control_panel.ATTR_CODE_ARM_REQUIRED: True, - ATTR_SUPPORTED_FEATURES: alarm_control_panel.const.SUPPORT_ALARM_TRIGGER - | alarm_control_panel.const.SUPPORT_ALARM_ARM_CUSTOM_BYPASS, + ATTR_SUPPORTED_FEATURES: AlarmControlPanelEntityFeature.TRIGGER + | AlarmControlPanelEntityFeature.ARM_CUSTOM_BYPASS, }, ), PIN_CONFIG, @@ -1662,7 +1690,9 @@ async def test_arm_disarm_disarm(hass: HomeAssistant) -> None: async def test_fan_speed(hass: HomeAssistant) -> None: """Test FanSpeed trait speed control support for fan domain.""" assert helpers.get_google_type(fan.DOMAIN, None) is not None - assert trait.FanSpeedTrait.supported(fan.DOMAIN, fan.SUPPORT_SET_SPEED, None, None) + assert trait.FanSpeedTrait.supported( + fan.DOMAIN, FanEntityFeature.SET_SPEED, None, None + ) trt = trait.FanSpeedTrait( hass, @@ -1700,7 +1730,9 @@ async def test_fan_speed(hass: HomeAssistant) -> None: async def test_fan_speed_without_percentage_step(hass: HomeAssistant) -> None: """Test FanSpeed trait speed control percentage step for fan domain.""" assert helpers.get_google_type(fan.DOMAIN, None) is not None - assert trait.FanSpeedTrait.supported(fan.DOMAIN, fan.SUPPORT_SET_SPEED, None, None) + assert trait.FanSpeedTrait.supported( + fan.DOMAIN, FanEntityFeature.SET_SPEED, None, None + ) trt = trait.FanSpeedTrait( hass, @@ -1787,7 +1819,9 @@ async def test_fan_speed_ordered( ): """Test FanSpeed trait speed control support for fan domain.""" assert helpers.get_google_type(fan.DOMAIN, None) is not None - assert trait.FanSpeedTrait.supported(fan.DOMAIN, fan.SUPPORT_SET_SPEED, None, None) + assert trait.FanSpeedTrait.supported( + fan.DOMAIN, FanEntityFeature.SET_SPEED, None, None + ) trt = trait.FanSpeedTrait( hass, @@ -1858,7 +1892,7 @@ async def test_fan_reverse( "percentage": 33, "percentage_step": 1.0, "direction": direction_state, - "supported_features": fan.SUPPORT_DIRECTION, + "supported_features": FanEntityFeature.DIRECTION, }, ), BASIC_CONFIG, @@ -1889,7 +1923,7 @@ async def test_climate_fan_speed(hass: HomeAssistant) -> None: """Test FanSpeed trait speed control support for climate domain.""" assert helpers.get_google_type(climate.DOMAIN, None) is not None assert trait.FanSpeedTrait.supported( - climate.DOMAIN, climate.SUPPORT_FAN_MODE, None, None + climate.DOMAIN, ClimateEntityFeature.FAN_MODE, None, None ) trt = trait.FanSpeedTrait( @@ -1951,7 +1985,7 @@ async def test_inputselector(hass: HomeAssistant) -> None: assert helpers.get_google_type(media_player.DOMAIN, None) is not None assert trait.InputSelectorTrait.supported( media_player.DOMAIN, - media_player.MediaPlayerEntityFeature.SELECT_SOURCE, + MediaPlayerEntityFeature.SELECT_SOURCE, None, None, ) @@ -2265,7 +2299,7 @@ async def test_modes_humidifier(hass: HomeAssistant) -> None: """Test Humidifier Mode trait.""" assert helpers.get_google_type(humidifier.DOMAIN, None) is not None assert trait.ModesTrait.supported( - humidifier.DOMAIN, humidifier.SUPPORT_MODES, None, None + humidifier.DOMAIN, HumidifierEntityFeature.MODES, None, None ) trt = trait.ModesTrait( @@ -2279,7 +2313,7 @@ async def test_modes_humidifier(hass: HomeAssistant) -> None: humidifier.MODE_AUTO, humidifier.MODE_AWAY, ], - ATTR_SUPPORTED_FEATURES: humidifier.SUPPORT_MODES, + ATTR_SUPPORTED_FEATURES: HumidifierEntityFeature.MODES, humidifier.ATTR_MIN_HUMIDITY: 30, humidifier.ATTR_MAX_HUMIDITY: 99, humidifier.ATTR_HUMIDITY: 50, @@ -2345,7 +2379,7 @@ async def test_sound_modes(hass: HomeAssistant) -> None: assert helpers.get_google_type(media_player.DOMAIN, None) is not None assert trait.ModesTrait.supported( media_player.DOMAIN, - media_player.MediaPlayerEntityFeature.SELECT_SOUND_MODE, + MediaPlayerEntityFeature.SELECT_SOUND_MODE, None, None, ) @@ -2420,7 +2454,9 @@ async def test_sound_modes(hass: HomeAssistant) -> None: async def test_preset_modes(hass: HomeAssistant) -> None: """Test Mode trait for fan preset modes.""" assert helpers.get_google_type(fan.DOMAIN, None) is not None - assert trait.ModesTrait.supported(fan.DOMAIN, fan.SUPPORT_PRESET_MODE, None, None) + assert trait.ModesTrait.supported( + fan.DOMAIN, FanEntityFeature.PRESET_MODE, None, None + ) trt = trait.ModesTrait( hass, @@ -2430,7 +2466,7 @@ async def test_preset_modes(hass: HomeAssistant) -> None: attributes={ fan.ATTR_PRESET_MODES: ["auto", "whoosh"], fan.ATTR_PRESET_MODE: "auto", - ATTR_SUPPORTED_FEATURES: fan.SUPPORT_PRESET_MODE, + ATTR_SUPPORTED_FEATURES: FanEntityFeature.PRESET_MODE, }, ), BASIC_CONFIG, @@ -2514,7 +2550,7 @@ async def test_openclose_cover(hass: HomeAssistant) -> None: """Test OpenClose trait support for cover domain.""" assert helpers.get_google_type(cover.DOMAIN, None) is not None assert trait.OpenCloseTrait.supported( - cover.DOMAIN, cover.SUPPORT_SET_POSITION, None, None + cover.DOMAIN, CoverEntityFeature.SET_POSITION, None, None ) trt = trait.OpenCloseTrait( @@ -2524,7 +2560,7 @@ async def test_openclose_cover(hass: HomeAssistant) -> None: cover.STATE_OPEN, { cover.ATTR_CURRENT_POSITION: 75, - ATTR_SUPPORTED_FEATURES: cover.SUPPORT_SET_POSITION, + ATTR_SUPPORTED_FEATURES: CoverEntityFeature.SET_POSITION, }, ), BASIC_CONFIG, @@ -2551,14 +2587,16 @@ async def test_openclose_cover_unknown_state(hass: HomeAssistant) -> None: """Test OpenClose trait support for cover domain with unknown state.""" assert helpers.get_google_type(cover.DOMAIN, None) is not None assert trait.OpenCloseTrait.supported( - cover.DOMAIN, cover.SUPPORT_SET_POSITION, None, None + cover.DOMAIN, CoverEntityFeature.SET_POSITION, None, None ) # No state trt = trait.OpenCloseTrait( hass, State( - "cover.bla", STATE_UNKNOWN, {ATTR_SUPPORTED_FEATURES: cover.SUPPORT_OPEN} + "cover.bla", + STATE_UNKNOWN, + {ATTR_SUPPORTED_FEATURES: CoverEntityFeature.OPEN}, ), BASIC_CONFIG, ) @@ -2581,7 +2619,7 @@ async def test_openclose_cover_assumed_state(hass: HomeAssistant) -> None: """Test OpenClose trait support for cover domain.""" assert helpers.get_google_type(cover.DOMAIN, None) is not None assert trait.OpenCloseTrait.supported( - cover.DOMAIN, cover.SUPPORT_SET_POSITION, None, None + cover.DOMAIN, CoverEntityFeature.SET_POSITION, None, None ) trt = trait.OpenCloseTrait( @@ -2591,7 +2629,7 @@ async def test_openclose_cover_assumed_state(hass: HomeAssistant) -> None: cover.STATE_OPEN, { ATTR_ASSUMED_STATE: True, - ATTR_SUPPORTED_FEATURES: cover.SUPPORT_SET_POSITION, + ATTR_SUPPORTED_FEATURES: CoverEntityFeature.SET_POSITION, }, ), BASIC_CONFIG, @@ -2634,14 +2672,17 @@ async def test_openclose_cover_no_position(hass: HomeAssistant) -> None: """Test OpenClose trait support for cover domain.""" assert helpers.get_google_type(cover.DOMAIN, None) is not None assert trait.OpenCloseTrait.supported( - cover.DOMAIN, cover.SUPPORT_OPEN | cover.SUPPORT_CLOSE, None, None + cover.DOMAIN, + CoverEntityFeature.OPEN | CoverEntityFeature.CLOSE, + None, + None, ) state = State( "cover.bla", cover.STATE_OPEN, { - ATTR_SUPPORTED_FEATURES: cover.SUPPORT_OPEN | cover.SUPPORT_CLOSE, + ATTR_SUPPORTED_FEATURES: CoverEntityFeature.OPEN | CoverEntityFeature.CLOSE, }, ) @@ -2695,10 +2736,10 @@ async def test_openclose_cover_secure(hass: HomeAssistant, device_class) -> None """Test OpenClose trait support for cover domain.""" assert helpers.get_google_type(cover.DOMAIN, device_class) is not None assert trait.OpenCloseTrait.supported( - cover.DOMAIN, cover.SUPPORT_SET_POSITION, device_class, None + cover.DOMAIN, CoverEntityFeature.SET_POSITION, device_class, None ) assert trait.OpenCloseTrait.might_2fa( - cover.DOMAIN, cover.SUPPORT_SET_POSITION, device_class + cover.DOMAIN, CoverEntityFeature.SET_POSITION, device_class ) trt = trait.OpenCloseTrait( @@ -2708,7 +2749,7 @@ async def test_openclose_cover_secure(hass: HomeAssistant, device_class) -> None cover.STATE_OPEN, { ATTR_DEVICE_CLASS: device_class, - ATTR_SUPPORTED_FEATURES: cover.SUPPORT_SET_POSITION, + ATTR_SUPPORTED_FEATURES: CoverEntityFeature.SET_POSITION, cover.ATTR_CURRENT_POSITION: 75, }, ), @@ -2796,7 +2837,7 @@ async def test_volume_media_player(hass: HomeAssistant) -> None: assert helpers.get_google_type(media_player.DOMAIN, None) is not None assert trait.VolumeTrait.supported( media_player.DOMAIN, - media_player.MediaPlayerEntityFeature.VOLUME_SET, + MediaPlayerEntityFeature.VOLUME_SET, None, None, ) @@ -2807,7 +2848,7 @@ async def test_volume_media_player(hass: HomeAssistant) -> None: "media_player.bla", media_player.STATE_PLAYING, { - ATTR_SUPPORTED_FEATURES: media_player.MediaPlayerEntityFeature.VOLUME_SET, + ATTR_SUPPORTED_FEATURES: MediaPlayerEntityFeature.VOLUME_SET, media_player.ATTR_MEDIA_VOLUME_LEVEL: 0.3, }, ), @@ -2850,7 +2891,7 @@ async def test_volume_media_player_relative(hass: HomeAssistant) -> None: """Test volume trait support for relative-volume-only media players.""" assert trait.VolumeTrait.supported( media_player.DOMAIN, - media_player.MediaPlayerEntityFeature.VOLUME_STEP, + MediaPlayerEntityFeature.VOLUME_STEP, None, None, ) @@ -2861,7 +2902,7 @@ async def test_volume_media_player_relative(hass: HomeAssistant) -> None: media_player.STATE_PLAYING, { ATTR_ASSUMED_STATE: True, - ATTR_SUPPORTED_FEATURES: media_player.MediaPlayerEntityFeature.VOLUME_STEP, + ATTR_SUPPORTED_FEATURES: MediaPlayerEntityFeature.VOLUME_STEP, }, ), BASIC_CONFIG, @@ -2918,8 +2959,7 @@ async def test_media_player_mute(hass: HomeAssistant) -> None: """Test volume trait support for muting.""" assert trait.VolumeTrait.supported( media_player.DOMAIN, - media_player.MediaPlayerEntityFeature.VOLUME_STEP - | media_player.MediaPlayerEntityFeature.VOLUME_MUTE, + MediaPlayerEntityFeature.VOLUME_STEP | MediaPlayerEntityFeature.VOLUME_MUTE, None, None, ) @@ -2930,8 +2970,8 @@ async def test_media_player_mute(hass: HomeAssistant) -> None: media_player.STATE_PLAYING, { ATTR_SUPPORTED_FEATURES: ( - media_player.MediaPlayerEntityFeature.VOLUME_STEP - | media_player.MediaPlayerEntityFeature.VOLUME_MUTE + MediaPlayerEntityFeature.VOLUME_STEP + | MediaPlayerEntityFeature.VOLUME_MUTE ), media_player.ATTR_MEDIA_VOLUME_MUTED: False, }, @@ -3095,8 +3135,8 @@ async def test_transport_control(hass: HomeAssistant) -> None: media_player.ATTR_MEDIA_POSITION_UPDATED_AT: now - timedelta(seconds=10), media_player.ATTR_MEDIA_VOLUME_LEVEL: 0.5, - ATTR_SUPPORTED_FEATURES: media_player.MediaPlayerEntityFeature.PLAY - | media_player.MediaPlayerEntityFeature.STOP, + ATTR_SUPPORTED_FEATURES: MediaPlayerEntityFeature.PLAY + | MediaPlayerEntityFeature.STOP, }, ), BASIC_CONFIG, @@ -3210,7 +3250,7 @@ async def test_media_state(hass: HomeAssistant, state) -> None: assert helpers.get_google_type(media_player.DOMAIN, None) is not None assert trait.TransportControlTrait.supported( - media_player.DOMAIN, media_player.MediaPlayerEntityFeature.PLAY, None, None + media_player.DOMAIN, MediaPlayerEntityFeature.PLAY, None, None ) trt = trait.MediaStateTrait( @@ -3222,8 +3262,8 @@ async def test_media_state(hass: HomeAssistant, state) -> None: media_player.ATTR_MEDIA_POSITION: 100, media_player.ATTR_MEDIA_DURATION: 200, media_player.ATTR_MEDIA_VOLUME_LEVEL: 0.5, - ATTR_SUPPORTED_FEATURES: media_player.MediaPlayerEntityFeature.PLAY - | media_player.MediaPlayerEntityFeature.STOP, + ATTR_SUPPORTED_FEATURES: MediaPlayerEntityFeature.PLAY + | MediaPlayerEntityFeature.STOP, }, ), BASIC_CONFIG, @@ -3244,14 +3284,14 @@ async def test_channel(hass: HomeAssistant) -> None: assert helpers.get_google_type(media_player.DOMAIN, None) is not None assert trait.ChannelTrait.supported( media_player.DOMAIN, - media_player.MediaPlayerEntityFeature.PLAY_MEDIA, + MediaPlayerEntityFeature.PLAY_MEDIA, media_player.MediaPlayerDeviceClass.TV, None, ) assert ( trait.ChannelTrait.supported( media_player.DOMAIN, - media_player.MediaPlayerEntityFeature.PLAY_MEDIA, + MediaPlayerEntityFeature.PLAY_MEDIA, None, None, )