Use VacuumEntityFeature in mqtt (#70570)

* Use VacuumEntityFeature in mqtt

* Fix tests
This commit is contained in:
epenet 2022-04-25 11:34:06 +02:00 committed by GitHub
parent 6988b0725d
commit a8a7359c6f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 64 additions and 82 deletions

View File

@ -6,18 +6,8 @@ import voluptuous as vol
from homeassistant.components.vacuum import ( from homeassistant.components.vacuum import (
ATTR_STATUS, ATTR_STATUS,
ENTITY_ID_FORMAT, ENTITY_ID_FORMAT,
SUPPORT_BATTERY,
SUPPORT_CLEAN_SPOT,
SUPPORT_FAN_SPEED,
SUPPORT_LOCATE,
SUPPORT_PAUSE,
SUPPORT_RETURN_HOME,
SUPPORT_SEND_COMMAND,
SUPPORT_STATUS,
SUPPORT_STOP,
SUPPORT_TURN_OFF,
SUPPORT_TURN_ON,
VacuumEntity, VacuumEntity,
VacuumEntityFeature,
) )
from homeassistant.const import ATTR_SUPPORTED_FEATURES, CONF_NAME from homeassistant.const import ATTR_SUPPORTED_FEATURES, CONF_NAME
from homeassistant.core import callback from homeassistant.core import callback
@ -33,36 +23,36 @@ from .const import MQTT_VACUUM_ATTRIBUTES_BLOCKED
from .schema import MQTT_VACUUM_SCHEMA, services_to_strings, strings_to_services from .schema import MQTT_VACUUM_SCHEMA, services_to_strings, strings_to_services
SERVICE_TO_STRING = { SERVICE_TO_STRING = {
SUPPORT_TURN_ON: "turn_on", VacuumEntityFeature.TURN_ON: "turn_on",
SUPPORT_TURN_OFF: "turn_off", VacuumEntityFeature.TURN_OFF: "turn_off",
SUPPORT_PAUSE: "pause", VacuumEntityFeature.PAUSE: "pause",
SUPPORT_STOP: "stop", VacuumEntityFeature.STOP: "stop",
SUPPORT_RETURN_HOME: "return_home", VacuumEntityFeature.RETURN_HOME: "return_home",
SUPPORT_FAN_SPEED: "fan_speed", VacuumEntityFeature.FAN_SPEED: "fan_speed",
SUPPORT_BATTERY: "battery", VacuumEntityFeature.BATTERY: "battery",
SUPPORT_STATUS: "status", VacuumEntityFeature.STATUS: "status",
SUPPORT_SEND_COMMAND: "send_command", VacuumEntityFeature.SEND_COMMAND: "send_command",
SUPPORT_LOCATE: "locate", VacuumEntityFeature.LOCATE: "locate",
SUPPORT_CLEAN_SPOT: "clean_spot", VacuumEntityFeature.CLEAN_SPOT: "clean_spot",
} }
STRING_TO_SERVICE = {v: k for k, v in SERVICE_TO_STRING.items()} STRING_TO_SERVICE = {v: k for k, v in SERVICE_TO_STRING.items()}
DEFAULT_SERVICES = ( DEFAULT_SERVICES = (
SUPPORT_TURN_ON VacuumEntityFeature.TURN_ON
| SUPPORT_TURN_OFF | VacuumEntityFeature.TURN_OFF
| SUPPORT_STOP | VacuumEntityFeature.STOP
| SUPPORT_RETURN_HOME | VacuumEntityFeature.RETURN_HOME
| SUPPORT_STATUS | VacuumEntityFeature.STATUS
| SUPPORT_BATTERY | VacuumEntityFeature.BATTERY
| SUPPORT_CLEAN_SPOT | VacuumEntityFeature.CLEAN_SPOT
) )
ALL_SERVICES = ( ALL_SERVICES = (
DEFAULT_SERVICES DEFAULT_SERVICES
| SUPPORT_PAUSE | VacuumEntityFeature.PAUSE
| SUPPORT_LOCATE | VacuumEntityFeature.LOCATE
| SUPPORT_FAN_SPEED | VacuumEntityFeature.FAN_SPEED
| SUPPORT_SEND_COMMAND | VacuumEntityFeature.SEND_COMMAND
) )
CONF_SUPPORTED_FEATURES = ATTR_SUPPORTED_FEATURES CONF_SUPPORTED_FEATURES = ATTR_SUPPORTED_FEATURES
@ -372,7 +362,7 @@ class MqttVacuum(MqttEntity, VacuumEntity):
def battery_icon(self): def battery_icon(self):
"""Return the battery icon for the vacuum cleaner. """Return the battery icon for the vacuum cleaner.
No need to check SUPPORT_BATTERY, this won't be called if battery_level is None. No need to check VacuumEntityFeature.BATTERY, this won't be called if battery_level is None.
""" """
return icon_for_battery_level( return icon_for_battery_level(
battery_level=self.battery_level, charging=self._charging battery_level=self.battery_level, charging=self._charging
@ -385,7 +375,7 @@ class MqttVacuum(MqttEntity, VacuumEntity):
async def async_turn_on(self, **kwargs): async def async_turn_on(self, **kwargs):
"""Turn the vacuum on.""" """Turn the vacuum on."""
if self.supported_features & SUPPORT_TURN_ON == 0: if self.supported_features & VacuumEntityFeature.TURN_ON == 0:
return return
await self.async_publish( await self.async_publish(
@ -400,7 +390,7 @@ class MqttVacuum(MqttEntity, VacuumEntity):
async def async_turn_off(self, **kwargs): async def async_turn_off(self, **kwargs):
"""Turn the vacuum off.""" """Turn the vacuum off."""
if self.supported_features & SUPPORT_TURN_OFF == 0: if self.supported_features & VacuumEntityFeature.TURN_OFF == 0:
return None return None
await self.async_publish( await self.async_publish(
@ -415,7 +405,7 @@ class MqttVacuum(MqttEntity, VacuumEntity):
async def async_stop(self, **kwargs): async def async_stop(self, **kwargs):
"""Stop the vacuum.""" """Stop the vacuum."""
if self.supported_features & SUPPORT_STOP == 0: if self.supported_features & VacuumEntityFeature.STOP == 0:
return None return None
await self.async_publish( await self.async_publish(
@ -430,7 +420,7 @@ class MqttVacuum(MqttEntity, VacuumEntity):
async def async_clean_spot(self, **kwargs): async def async_clean_spot(self, **kwargs):
"""Perform a spot clean-up.""" """Perform a spot clean-up."""
if self.supported_features & SUPPORT_CLEAN_SPOT == 0: if self.supported_features & VacuumEntityFeature.CLEAN_SPOT == 0:
return None return None
await self.async_publish( await self.async_publish(
@ -445,7 +435,7 @@ class MqttVacuum(MqttEntity, VacuumEntity):
async def async_locate(self, **kwargs): async def async_locate(self, **kwargs):
"""Locate the vacuum (usually by playing a song).""" """Locate the vacuum (usually by playing a song)."""
if self.supported_features & SUPPORT_LOCATE == 0: if self.supported_features & VacuumEntityFeature.LOCATE == 0:
return None return None
await self.async_publish( await self.async_publish(
@ -460,7 +450,7 @@ class MqttVacuum(MqttEntity, VacuumEntity):
async def async_start_pause(self, **kwargs): async def async_start_pause(self, **kwargs):
"""Start, pause or resume the cleaning task.""" """Start, pause or resume the cleaning task."""
if self.supported_features & SUPPORT_PAUSE == 0: if self.supported_features & VacuumEntityFeature.PAUSE == 0:
return None return None
await self.async_publish( await self.async_publish(
@ -475,7 +465,7 @@ class MqttVacuum(MqttEntity, VacuumEntity):
async def async_return_to_base(self, **kwargs): async def async_return_to_base(self, **kwargs):
"""Tell the vacuum to return to its dock.""" """Tell the vacuum to return to its dock."""
if self.supported_features & SUPPORT_RETURN_HOME == 0: if self.supported_features & VacuumEntityFeature.RETURN_HOME == 0:
return None return None
await self.async_publish( await self.async_publish(
@ -491,7 +481,7 @@ class MqttVacuum(MqttEntity, VacuumEntity):
async def async_set_fan_speed(self, fan_speed, **kwargs): async def async_set_fan_speed(self, fan_speed, **kwargs):
"""Set fan speed.""" """Set fan speed."""
if ( if (
self.supported_features & SUPPORT_FAN_SPEED == 0 self.supported_features & VacuumEntityFeature.FAN_SPEED == 0
) or fan_speed not in self._fan_speed_list: ) or fan_speed not in self._fan_speed_list:
return None return None
@ -507,7 +497,7 @@ class MqttVacuum(MqttEntity, VacuumEntity):
async def async_send_command(self, command, params=None, **kwargs): async def async_send_command(self, command, params=None, **kwargs):
"""Send a command to a vacuum cleaner.""" """Send a command to a vacuum cleaner."""
if self.supported_features & SUPPORT_SEND_COMMAND == 0: if self.supported_features & VacuumEntityFeature.SEND_COMMAND == 0:
return return
if params: if params:
message = {"command": command} message = {"command": command}

View File

@ -9,17 +9,8 @@ from homeassistant.components.vacuum import (
STATE_DOCKED, STATE_DOCKED,
STATE_ERROR, STATE_ERROR,
STATE_RETURNING, STATE_RETURNING,
SUPPORT_BATTERY,
SUPPORT_CLEAN_SPOT,
SUPPORT_FAN_SPEED,
SUPPORT_LOCATE,
SUPPORT_PAUSE,
SUPPORT_RETURN_HOME,
SUPPORT_SEND_COMMAND,
SUPPORT_START,
SUPPORT_STATUS,
SUPPORT_STOP,
StateVacuumEntity, StateVacuumEntity,
VacuumEntityFeature,
) )
from homeassistant.const import ( from homeassistant.const import (
ATTR_SUPPORTED_FEATURES, ATTR_SUPPORTED_FEATURES,
@ -45,35 +36,35 @@ from .const import MQTT_VACUUM_ATTRIBUTES_BLOCKED
from .schema import MQTT_VACUUM_SCHEMA, services_to_strings, strings_to_services from .schema import MQTT_VACUUM_SCHEMA, services_to_strings, strings_to_services
SERVICE_TO_STRING = { SERVICE_TO_STRING = {
SUPPORT_START: "start", VacuumEntityFeature.START: "start",
SUPPORT_PAUSE: "pause", VacuumEntityFeature.PAUSE: "pause",
SUPPORT_STOP: "stop", VacuumEntityFeature.STOP: "stop",
SUPPORT_RETURN_HOME: "return_home", VacuumEntityFeature.RETURN_HOME: "return_home",
SUPPORT_FAN_SPEED: "fan_speed", VacuumEntityFeature.FAN_SPEED: "fan_speed",
SUPPORT_BATTERY: "battery", VacuumEntityFeature.BATTERY: "battery",
SUPPORT_STATUS: "status", VacuumEntityFeature.STATUS: "status",
SUPPORT_SEND_COMMAND: "send_command", VacuumEntityFeature.SEND_COMMAND: "send_command",
SUPPORT_LOCATE: "locate", VacuumEntityFeature.LOCATE: "locate",
SUPPORT_CLEAN_SPOT: "clean_spot", VacuumEntityFeature.CLEAN_SPOT: "clean_spot",
} }
STRING_TO_SERVICE = {v: k for k, v in SERVICE_TO_STRING.items()} STRING_TO_SERVICE = {v: k for k, v in SERVICE_TO_STRING.items()}
DEFAULT_SERVICES = ( DEFAULT_SERVICES = (
SUPPORT_START VacuumEntityFeature.START
| SUPPORT_STOP | VacuumEntityFeature.STOP
| SUPPORT_RETURN_HOME | VacuumEntityFeature.RETURN_HOME
| SUPPORT_STATUS | VacuumEntityFeature.STATUS
| SUPPORT_BATTERY | VacuumEntityFeature.BATTERY
| SUPPORT_CLEAN_SPOT | VacuumEntityFeature.CLEAN_SPOT
) )
ALL_SERVICES = ( ALL_SERVICES = (
DEFAULT_SERVICES DEFAULT_SERVICES
| SUPPORT_PAUSE | VacuumEntityFeature.PAUSE
| SUPPORT_LOCATE | VacuumEntityFeature.LOCATE
| SUPPORT_FAN_SPEED | VacuumEntityFeature.FAN_SPEED
| SUPPORT_SEND_COMMAND | VacuumEntityFeature.SEND_COMMAND
) )
BATTERY = "battery_level" BATTERY = "battery_level"
@ -258,7 +249,7 @@ class MqttStateVacuum(MqttEntity, StateVacuumEntity):
async def async_start(self): async def async_start(self):
"""Start the vacuum.""" """Start the vacuum."""
if self.supported_features & SUPPORT_START == 0: if self.supported_features & VacuumEntityFeature.START == 0:
return None return None
await self.async_publish( await self.async_publish(
self._command_topic, self._command_topic,
@ -270,7 +261,7 @@ class MqttStateVacuum(MqttEntity, StateVacuumEntity):
async def async_pause(self): async def async_pause(self):
"""Pause the vacuum.""" """Pause the vacuum."""
if self.supported_features & SUPPORT_PAUSE == 0: if self.supported_features & VacuumEntityFeature.PAUSE == 0:
return None return None
await self.async_publish( await self.async_publish(
self._command_topic, self._command_topic,
@ -282,7 +273,7 @@ class MqttStateVacuum(MqttEntity, StateVacuumEntity):
async def async_stop(self, **kwargs): async def async_stop(self, **kwargs):
"""Stop the vacuum.""" """Stop the vacuum."""
if self.supported_features & SUPPORT_STOP == 0: if self.supported_features & VacuumEntityFeature.STOP == 0:
return None return None
await self.async_publish( await self.async_publish(
self._command_topic, self._command_topic,
@ -294,7 +285,7 @@ class MqttStateVacuum(MqttEntity, StateVacuumEntity):
async def async_set_fan_speed(self, fan_speed, **kwargs): async def async_set_fan_speed(self, fan_speed, **kwargs):
"""Set fan speed.""" """Set fan speed."""
if (self.supported_features & SUPPORT_FAN_SPEED == 0) or ( if (self.supported_features & VacuumEntityFeature.FAN_SPEED == 0) or (
fan_speed not in self._fan_speed_list fan_speed not in self._fan_speed_list
): ):
return None return None
@ -308,7 +299,7 @@ class MqttStateVacuum(MqttEntity, StateVacuumEntity):
async def async_return_to_base(self, **kwargs): async def async_return_to_base(self, **kwargs):
"""Tell the vacuum to return to its dock.""" """Tell the vacuum to return to its dock."""
if self.supported_features & SUPPORT_RETURN_HOME == 0: if self.supported_features & VacuumEntityFeature.RETURN_HOME == 0:
return None return None
await self.async_publish( await self.async_publish(
self._command_topic, self._command_topic,
@ -320,7 +311,7 @@ class MqttStateVacuum(MqttEntity, StateVacuumEntity):
async def async_clean_spot(self, **kwargs): async def async_clean_spot(self, **kwargs):
"""Perform a spot clean-up.""" """Perform a spot clean-up."""
if self.supported_features & SUPPORT_CLEAN_SPOT == 0: if self.supported_features & VacuumEntityFeature.CLEAN_SPOT == 0:
return None return None
await self.async_publish( await self.async_publish(
self._command_topic, self._command_topic,
@ -332,7 +323,7 @@ class MqttStateVacuum(MqttEntity, StateVacuumEntity):
async def async_locate(self, **kwargs): async def async_locate(self, **kwargs):
"""Locate the vacuum (usually by playing a song).""" """Locate the vacuum (usually by playing a song)."""
if self.supported_features & SUPPORT_LOCATE == 0: if self.supported_features & VacuumEntityFeature.LOCATE == 0:
return None return None
await self.async_publish( await self.async_publish(
self._command_topic, self._command_topic,
@ -344,7 +335,7 @@ class MqttStateVacuum(MqttEntity, StateVacuumEntity):
async def async_send_command(self, command, params=None, **kwargs): async def async_send_command(self, command, params=None, **kwargs):
"""Send a command to a vacuum cleaner.""" """Send a command to a vacuum cleaner."""
if self.supported_features & SUPPORT_SEND_COMMAND == 0: if self.supported_features & VacuumEntityFeature.SEND_COMMAND == 0:
return None return None
if params: if params:
message = {"command": command} message = {"command": command}

View File

@ -27,6 +27,7 @@ from homeassistant.components.vacuum import (
ATTR_FAN_SPEED, ATTR_FAN_SPEED,
ATTR_FAN_SPEED_LIST, ATTR_FAN_SPEED_LIST,
ATTR_STATUS, ATTR_STATUS,
VacuumEntityFeature,
) )
from homeassistant.const import CONF_NAME, CONF_PLATFORM, STATE_OFF, STATE_ON from homeassistant.const import CONF_NAME, CONF_PLATFORM, STATE_OFF, STATE_ON
from homeassistant.setup import async_setup_component from homeassistant.setup import async_setup_component
@ -412,7 +413,7 @@ async def test_status_no_fan_speed_list(hass, mqtt_mock):
If the vacuum doesn't support fan speed, fan speed list should be None. If the vacuum doesn't support fan speed, fan speed list should be None.
""" """
config = deepcopy(DEFAULT_CONFIG) config = deepcopy(DEFAULT_CONFIG)
services = ALL_SERVICES - mqttvacuum.SUPPORT_FAN_SPEED services = ALL_SERVICES - VacuumEntityFeature.FAN_SPEED
config[mqttvacuum.CONF_SUPPORTED_FEATURES] = services_to_strings( config[mqttvacuum.CONF_SUPPORTED_FEATURES] = services_to_strings(
services, SERVICE_TO_STRING services, SERVICE_TO_STRING
) )