MQTT Vacuum State Device (#23171)

* add StateVacuum MQTT

homeassistant/components/mqtt/vacuum/__init__.py

@@ -0,0 +1,97 @@
+"""
+Support for MQTT vacuums.
+
+For more details about this platform, please refer to the documentation at
+https://www.home-assistant.io/components/vacuum.mqtt/
+"""
+import logging
+
+import voluptuous as vol
+
+from homeassistant.components.vacuum import DOMAIN
+from homeassistant.components.mqtt import ATTR_DISCOVERY_HASH
+from homeassistant.components.mqtt.discovery import (
+    MQTT_DISCOVERY_NEW, clear_discovery_hash)
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
+
+_LOGGER = logging.getLogger(__name__)
+
+CONF_SCHEMA = 'schema'
+LEGACY = 'legacy'
+STATE = 'state'
+
+
+def validate_mqtt_vacuum(value):
+    """Validate MQTT vacuum schema."""
+    from . import schema_legacy
+    from . import schema_state
+
+    schemas = {
+        LEGACY: schema_legacy.PLATFORM_SCHEMA_LEGACY,
+        STATE: schema_state.PLATFORM_SCHEMA_STATE,
+    }
+    return schemas[value[CONF_SCHEMA]](value)
+
+
+def services_to_strings(services, service_to_string):
+    """Convert SUPPORT_* service bitmask to list of service strings."""
+    strings = []
+    for service in service_to_string:
+        if service & services:
+            strings.append(service_to_string[service])
+    return strings
+
+
+def strings_to_services(strings, string_to_service):
+    """Convert service strings to SUPPORT_* service bitmask."""
+    services = 0
+    for string in strings:
+        services |= string_to_service[string]
+    return services
+
+
+MQTT_VACUUM_SCHEMA = vol.Schema({
+    vol.Optional(CONF_SCHEMA, default=LEGACY): vol.All(
+        vol.Lower, vol.Any(LEGACY, STATE))
+})
+
+PLATFORM_SCHEMA = vol.All(MQTT_VACUUM_SCHEMA.extend({
+}, extra=vol.ALLOW_EXTRA), validate_mqtt_vacuum)
+
+
+async def async_setup_platform(hass, config, async_add_entities,
+                               discovery_info=None):
+    """Set up MQTT vacuum through configuration.yaml."""
+    await _async_setup_entity(config, async_add_entities,
+                              discovery_info)
+
+
+async def async_setup_entry(hass, config_entry, async_add_entities):
+    """Set up MQTT vacuum dynamically through MQTT discovery."""
+    async def async_discover(discovery_payload):
+        """Discover and add a MQTT vacuum."""
+        try:
+            discovery_hash = discovery_payload.pop(ATTR_DISCOVERY_HASH)
+            config = PLATFORM_SCHEMA(discovery_payload)
+            await _async_setup_entity(config, async_add_entities, config_entry,
+                                      discovery_hash)
+        except Exception:
+            if discovery_hash:
+                clear_discovery_hash(hass, discovery_hash)
+            raise
+
+    async_dispatcher_connect(
+        hass, MQTT_DISCOVERY_NEW.format(DOMAIN, 'mqtt'), async_discover)
+
+
+async def _async_setup_entity(config, async_add_entities, config_entry,
+                              discovery_hash=None):
+    """Set up the MQTT vacuum."""
+    from . import schema_legacy
+    from . import schema_state
+    setup_entity = {
+        LEGACY: schema_legacy.async_setup_entity_legacy,
+        STATE: schema_state.async_setup_entity_state,
+    }
+    await setup_entity[config[CONF_SCHEMA]](
+        config, async_add_entities, config_entry, discovery_hash)

homeassistant/components/mqtt/vacuum/schema_legacy.py

@@ -1,4 +1,4 @@
-"""Support for a generic MQTT vacuum."""
+"""Support for Legacy MQTT vacuum."""
 import logging
 import json
 
@@ -6,20 +6,20 @@ import voluptuous as vol
 
 from homeassistant.components import mqtt
 from homeassistant.components.vacuum import (
-    DOMAIN, SUPPORT_BATTERY, SUPPORT_CLEAN_SPOT, SUPPORT_FAN_SPEED,
+    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,
     VacuumDevice)
 from homeassistant.const import ATTR_SUPPORTED_FEATURES, CONF_DEVICE, CONF_NAME
 from homeassistant.core import callback
 import homeassistant.helpers.config_validation as cv
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.icon import icon_for_battery_level
 
-from . import (
+from homeassistant.components.mqtt import (
-    ATTR_DISCOVERY_HASH, CONF_UNIQUE_ID, MqttAttributes, MqttAvailability,
+    CONF_UNIQUE_ID, MqttAttributes, MqttAvailability,
     MqttDiscoveryUpdate, MqttEntityDeviceInfo, subscription)
-from .discovery import MQTT_DISCOVERY_NEW, clear_discovery_hash
+
+from . import MQTT_VACUUM_SCHEMA, services_to_strings, strings_to_services
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -39,24 +39,6 @@ SERVICE_TO_STRING = {
 
 STRING_TO_SERVICE = {v: k for k, v in SERVICE_TO_STRING.items()}
 
-
-def services_to_strings(services):
-    """Convert SUPPORT_* service bitmask to list of service strings."""
-    strings = []
-    for service in SERVICE_TO_STRING:
-        if service & services:
-            strings.append(SERVICE_TO_STRING[service])
-    return strings
-
-
-def strings_to_services(strings):
-    """Convert service strings to SUPPORT_* service bitmask."""
-    services = 0
-    for string in strings:
-        services |= STRING_TO_SERVICE[string]
-    return services
-
-
 DEFAULT_SERVICES = SUPPORT_TURN_ON | SUPPORT_TURN_OFF | SUPPORT_STOP |\
                    SUPPORT_RETURN_HOME | SUPPORT_STATUS | SUPPORT_BATTERY |\
                    SUPPORT_CLEAN_SPOT
@@ -96,9 +78,10 @@ DEFAULT_PAYLOAD_STOP = 'stop'
 DEFAULT_PAYLOAD_TURN_OFF = 'turn_off'
 DEFAULT_PAYLOAD_TURN_ON = 'turn_on'
 DEFAULT_RETAIN = False
-DEFAULT_SERVICE_STRINGS = services_to_strings(DEFAULT_SERVICES)
+DEFAULT_SERVICE_STRINGS = services_to_strings(
+    DEFAULT_SERVICES, SERVICE_TO_STRING)
 
-PLATFORM_SCHEMA = mqtt.MQTT_BASE_PLATFORM_SCHEMA.extend({
+PLATFORM_SCHEMA_LEGACY = mqtt.MQTT_BASE_PLATFORM_SCHEMA.extend({
     vol.Inclusive(CONF_BATTERY_LEVEL_TEMPLATE, 'battery'): cv.template,
     vol.Inclusive(CONF_BATTERY_LEVEL_TOPIC,
                   'battery'): mqtt.valid_publish_topic,
@@ -137,44 +120,19 @@ PLATFORM_SCHEMA = mqtt.MQTT_BASE_PLATFORM_SCHEMA.extend({
     vol.Optional(mqtt.CONF_COMMAND_TOPIC): mqtt.valid_publish_topic,
     vol.Optional(mqtt.CONF_RETAIN, default=DEFAULT_RETAIN): cv.boolean,
 }).extend(mqtt.MQTT_AVAILABILITY_SCHEMA.schema).extend(
-    mqtt.MQTT_JSON_ATTRS_SCHEMA.schema)
+    mqtt.MQTT_JSON_ATTRS_SCHEMA.schema).extend(MQTT_VACUUM_SCHEMA.schema)
 
 
-async def async_setup_platform(hass, config, async_add_entities,
+async def async_setup_entity_legacy(config, async_add_entities,
-                               discovery_info=None):
+                                    config_entry, discovery_hash):
-    """Set up MQTT vacuum through configuration.yaml."""
+    """Set up a MQTT Vacuum Legacy."""
-    await _async_setup_entity(config, async_add_entities,
-                              discovery_info)
-
-
-async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up MQTT vacuum dynamically through MQTT discovery."""
-    async def async_discover(discovery_payload):
-        """Discover and add a MQTT vacuum."""
-        try:
-            discovery_hash = discovery_payload.pop(ATTR_DISCOVERY_HASH)
-            config = PLATFORM_SCHEMA(discovery_payload)
-            await _async_setup_entity(config, async_add_entities, config_entry,
-                                      discovery_hash)
-        except Exception:
-            if discovery_hash:
-                clear_discovery_hash(hass, discovery_hash)
-            raise
-
-    async_dispatcher_connect(
-        hass, MQTT_DISCOVERY_NEW.format(DOMAIN, 'mqtt'), async_discover)
-
-
-async def _async_setup_entity(config, async_add_entities, config_entry,
-                              discovery_hash=None):
-    """Set up the MQTT vacuum."""
     async_add_entities([MqttVacuum(config, config_entry, discovery_hash)])
 
 
 # pylint: disable=too-many-ancestors
 class MqttVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
                  MqttEntityDeviceInfo, VacuumDevice):
-    """Representation of a MQTT-controlled vacuum."""
+    """Representation of a MQTT-controlled legacy vacuum."""
 
     def __init__(self, config, config_entry, discovery_info):
         """Initialize the vacuum."""
@@ -204,7 +162,7 @@ class MqttVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
         self._name = config[CONF_NAME]
         supported_feature_strings = config[CONF_SUPPORTED_FEATURES]
         self._supported_features = strings_to_services(
-            supported_feature_strings
+            supported_feature_strings, STRING_TO_SERVICE
         )
         self._fan_speed_list = config[CONF_FAN_SPEED_LIST]
         self._qos = config[mqtt.CONF_QOS]
@@ -248,7 +206,7 @@ class MqttVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
 
     async def discovery_update(self, discovery_payload):
         """Handle updated discovery message."""
-        config = PLATFORM_SCHEMA(discovery_payload)
+        config = PLATFORM_SCHEMA_LEGACY(discovery_payload)
         self._setup_from_config(config)
         await self.attributes_discovery_update(config)
         await self.availability_discovery_update(config)
@@ -374,7 +332,7 @@ class MqttVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
     def status(self):
         """Return a status string for the vacuum."""
         if self.supported_features & SUPPORT_STATUS == 0:
-            return
+            return None
 
         return self._status
 
@@ -382,7 +340,7 @@ class MqttVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
     def fan_speed(self):
         """Return the status of the vacuum."""
         if self.supported_features & SUPPORT_FAN_SPEED == 0:
-            return
+            return None
 
         return self._fan_speed
 
@@ -429,7 +387,7 @@ class MqttVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
     async def async_turn_off(self, **kwargs):
         """Turn the vacuum off."""
         if self.supported_features & SUPPORT_TURN_OFF == 0:
-            return
+            return None
 
         mqtt.async_publish(self.hass, self._command_topic,
                            self._payloads[CONF_PAYLOAD_TURN_OFF],
@@ -440,7 +398,7 @@ class MqttVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
     async def async_stop(self, **kwargs):
         """Stop the vacuum."""
         if self.supported_features & SUPPORT_STOP == 0:
-            return
+            return None
 
         mqtt.async_publish(self.hass, self._command_topic,
                            self._payloads[CONF_PAYLOAD_STOP],
@@ -451,7 +409,7 @@ class MqttVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
     async def async_clean_spot(self, **kwargs):
         """Perform a spot clean-up."""
         if self.supported_features & SUPPORT_CLEAN_SPOT == 0:
-            return
+            return None
 
         mqtt.async_publish(self.hass, self._command_topic,
                            self._payloads[CONF_PAYLOAD_CLEAN_SPOT],
@@ -462,7 +420,7 @@ class MqttVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
     async def async_locate(self, **kwargs):
         """Locate the vacuum (usually by playing a song)."""
         if self.supported_features & SUPPORT_LOCATE == 0:
-            return
+            return None
 
         mqtt.async_publish(self.hass, self._command_topic,
                            self._payloads[CONF_PAYLOAD_LOCATE],
@@ -473,7 +431,7 @@ class MqttVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
     async def async_start_pause(self, **kwargs):
         """Start, pause or resume the cleaning task."""
         if self.supported_features & SUPPORT_PAUSE == 0:
-            return
+            return None
 
         mqtt.async_publish(self.hass, self._command_topic,
                            self._payloads[CONF_PAYLOAD_START_PAUSE],
@@ -484,7 +442,7 @@ class MqttVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
     async def async_return_to_base(self, **kwargs):
         """Tell the vacuum to return to its dock."""
         if self.supported_features & SUPPORT_RETURN_HOME == 0:
-            return
+            return None
 
         mqtt.async_publish(self.hass, self._command_topic,
                            self._payloads[CONF_PAYLOAD_RETURN_TO_BASE],
@@ -494,10 +452,9 @@ class MqttVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
 
     async def async_set_fan_speed(self, fan_speed, **kwargs):
         """Set fan speed."""
-        if self.supported_features & SUPPORT_FAN_SPEED == 0:
+        if ((self.supported_features & SUPPORT_FAN_SPEED == 0) or
-            return
+                fan_speed not in self._fan_speed_list):
-        if not self._fan_speed_list or fan_speed not in self._fan_speed_list:
+            return None
-            return
 
         mqtt.async_publish(self.hass, self._set_fan_speed_topic,
                            fan_speed, self._qos, self._retain)

homeassistant/components/mqtt/vacuum/schema_state.py

@@ -0,0 +1,339 @@
+"""Support for a State MQTT vacuum."""
+import logging
+import json
+
+import voluptuous as vol
+
+from homeassistant.components import mqtt
+from homeassistant.components.vacuum import (
+    SUPPORT_BATTERY, SUPPORT_CLEAN_SPOT, SUPPORT_FAN_SPEED, SUPPORT_START,
+    SUPPORT_LOCATE, SUPPORT_PAUSE, SUPPORT_RETURN_HOME, SUPPORT_SEND_COMMAND,
+    SUPPORT_STATUS, SUPPORT_STOP, STATE_CLEANING, STATE_DOCKED, STATE_PAUSED,
+    STATE_IDLE, STATE_RETURNING, STATE_ERROR, StateVacuumDevice)
+from homeassistant.const import (
+    ATTR_SUPPORTED_FEATURES, CONF_DEVICE, CONF_NAME, CONF_VALUE_TEMPLATE)
+from homeassistant.core import callback
+import homeassistant.helpers.config_validation as cv
+
+from homeassistant.components.mqtt import (
+    CONF_UNIQUE_ID, MqttAttributes, MqttAvailability,
+    MqttDiscoveryUpdate, MqttEntityDeviceInfo, subscription,
+    CONF_COMMAND_TOPIC, CONF_RETAIN, CONF_STATE_TOPIC, CONF_QOS)
+
+from . import MQTT_VACUUM_SCHEMA, services_to_strings, strings_to_services
+
+_LOGGER = logging.getLogger(__name__)
+
+SERVICE_TO_STRING = {
+    SUPPORT_START: 'start',
+    SUPPORT_PAUSE: 'pause',
+    SUPPORT_STOP: 'stop',
+    SUPPORT_RETURN_HOME: 'return_home',
+    SUPPORT_FAN_SPEED: 'fan_speed',
+    SUPPORT_BATTERY: 'battery',
+    SUPPORT_STATUS: 'status',
+    SUPPORT_SEND_COMMAND: 'send_command',
+    SUPPORT_LOCATE: 'locate',
+    SUPPORT_CLEAN_SPOT: 'clean_spot',
+}
+
+STRING_TO_SERVICE = {v: k for k, v in SERVICE_TO_STRING.items()}
+
+
+DEFAULT_SERVICES = SUPPORT_START | SUPPORT_STOP |\
+                   SUPPORT_RETURN_HOME | SUPPORT_STATUS | SUPPORT_BATTERY |\
+                   SUPPORT_CLEAN_SPOT
+ALL_SERVICES = DEFAULT_SERVICES | SUPPORT_PAUSE | SUPPORT_LOCATE |\
+               SUPPORT_FAN_SPEED | SUPPORT_SEND_COMMAND
+
+BATTERY = 'battery_level'
+FAN_SPEED = 'fan_speed'
+STATE = "state"
+
+POSSIBLE_STATES = {
+    STATE_IDLE: STATE_IDLE,
+    STATE_DOCKED: STATE_DOCKED,
+    STATE_ERROR: STATE_ERROR,
+    STATE_PAUSED: STATE_PAUSED,
+    STATE_RETURNING: STATE_RETURNING,
+    STATE_CLEANING: STATE_CLEANING,
+}
+
+CONF_SUPPORTED_FEATURES = ATTR_SUPPORTED_FEATURES
+CONF_PAYLOAD_TURN_ON = 'payload_turn_on'
+CONF_PAYLOAD_TURN_OFF = 'payload_turn_off'
+CONF_PAYLOAD_RETURN_TO_BASE = 'payload_return_to_base'
+CONF_PAYLOAD_STOP = 'payload_stop'
+CONF_PAYLOAD_CLEAN_SPOT = 'payload_clean_spot'
+CONF_PAYLOAD_LOCATE = 'payload_locate'
+CONF_PAYLOAD_START = 'payload_start'
+CONF_PAYLOAD_PAUSE = 'payload_pause'
+CONF_STATE_TEMPLATE = 'state_template'
+CONF_SET_FAN_SPEED_TOPIC = 'set_fan_speed_topic'
+CONF_FAN_SPEED_LIST = 'fan_speed_list'
+CONF_SEND_COMMAND_TOPIC = 'send_command_topic'
+
+DEFAULT_NAME = 'MQTT State Vacuum'
+DEFAULT_RETAIN = False
+DEFAULT_SERVICE_STRINGS = services_to_strings(
+    DEFAULT_SERVICES, SERVICE_TO_STRING)
+DEFAULT_PAYLOAD_RETURN_TO_BASE = 'return_to_base'
+DEFAULT_PAYLOAD_STOP = 'stop'
+DEFAULT_PAYLOAD_CLEAN_SPOT = 'clean_spot'
+DEFAULT_PAYLOAD_LOCATE = 'locate'
+DEFAULT_PAYLOAD_START = 'start'
+DEFAULT_PAYLOAD_PAUSE = 'pause'
+
+PLATFORM_SCHEMA_STATE = mqtt.MQTT_BASE_PLATFORM_SCHEMA.extend({
+    vol.Optional(CONF_DEVICE): mqtt.MQTT_ENTITY_DEVICE_INFO_SCHEMA,
+    vol.Optional(CONF_FAN_SPEED_LIST, default=[]):
+        vol.All(cv.ensure_list, [cv.string]),
+    vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
+    vol.Optional(CONF_PAYLOAD_CLEAN_SPOT,
+                 default=DEFAULT_PAYLOAD_CLEAN_SPOT): cv.string,
+    vol.Optional(CONF_PAYLOAD_LOCATE,
+                 default=DEFAULT_PAYLOAD_LOCATE): cv.string,
+    vol.Optional(CONF_PAYLOAD_RETURN_TO_BASE,
+                 default=DEFAULT_PAYLOAD_RETURN_TO_BASE): cv.string,
+    vol.Optional(CONF_PAYLOAD_START,
+                 default=DEFAULT_PAYLOAD_START): cv.string,
+    vol.Optional(CONF_PAYLOAD_PAUSE,
+                 default=DEFAULT_PAYLOAD_PAUSE): cv.string,
+    vol.Optional(CONF_PAYLOAD_STOP, default=DEFAULT_PAYLOAD_STOP): cv.string,
+    vol.Optional(CONF_SEND_COMMAND_TOPIC): mqtt.valid_publish_topic,
+    vol.Optional(CONF_SET_FAN_SPEED_TOPIC): mqtt.valid_publish_topic,
+    vol.Optional(CONF_VALUE_TEMPLATE): cv.template,
+    vol.Optional(CONF_STATE_TOPIC): mqtt.valid_publish_topic,
+    vol.Optional(CONF_SUPPORTED_FEATURES, default=DEFAULT_SERVICE_STRINGS):
+        vol.All(cv.ensure_list, [vol.In(STRING_TO_SERVICE.keys())]),
+    vol.Optional(CONF_UNIQUE_ID): cv.string,
+    vol.Optional(CONF_COMMAND_TOPIC): mqtt.valid_publish_topic,
+    vol.Optional(CONF_RETAIN, default=DEFAULT_RETAIN): cv.boolean,
+}).extend(mqtt.MQTT_AVAILABILITY_SCHEMA.schema).extend(
+    mqtt.MQTT_JSON_ATTRS_SCHEMA.schema).extend(MQTT_VACUUM_SCHEMA.schema)
+
+
+async def async_setup_entity_state(config, async_add_entities,
+                                   config_entry, discovery_hash):
+    """Set up a State MQTT Vacuum."""
+    async_add_entities([MqttStateVacuum(config, config_entry, discovery_hash)])
+
+
+# pylint: disable=too-many-ancestors
+class MqttStateVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
+                      MqttEntityDeviceInfo, StateVacuumDevice):
+    """Representation of a MQTT-controlled state vacuum."""
+
+    def __init__(self, config, config_entry, discovery_info):
+        """Initialize the vacuum."""
+        self._state = None
+        self._state_attrs = {}
+        self._fan_speed_list = []
+        self._sub_state = None
+        self._unique_id = config.get(CONF_UNIQUE_ID)
+
+        # Load config
+        self._setup_from_config(config)
+
+        device_config = config.get(CONF_DEVICE)
+
+        MqttAttributes.__init__(self, config)
+        MqttAvailability.__init__(self, config)
+        MqttDiscoveryUpdate.__init__(self, discovery_info,
+                                     self.discovery_update)
+        MqttEntityDeviceInfo.__init__(self, device_config, config_entry)
+
+    def _setup_from_config(self, config):
+        self._config = config
+        self._name = config[CONF_NAME]
+        supported_feature_strings = config[CONF_SUPPORTED_FEATURES]
+        self._supported_features = strings_to_services(
+            supported_feature_strings, STRING_TO_SERVICE
+        )
+        self._fan_speed_list = config[CONF_FAN_SPEED_LIST]
+        self._command_topic = config.get(mqtt.CONF_COMMAND_TOPIC)
+        self._set_fan_speed_topic = config.get(CONF_SET_FAN_SPEED_TOPIC)
+        self._send_command_topic = config.get(CONF_SEND_COMMAND_TOPIC)
+
+        self._payloads = {
+            key: config.get(key) for key in (
+                CONF_PAYLOAD_START,
+                CONF_PAYLOAD_PAUSE,
+                CONF_PAYLOAD_STOP,
+                CONF_PAYLOAD_RETURN_TO_BASE,
+                CONF_PAYLOAD_CLEAN_SPOT,
+                CONF_PAYLOAD_LOCATE
+            )
+        }
+
+    async def discovery_update(self, discovery_payload):
+        """Handle updated discovery message."""
+        config = PLATFORM_SCHEMA_STATE(discovery_payload)
+        self._setup_from_config(config)
+        await self.attributes_discovery_update(config)
+        await self.availability_discovery_update(config)
+        await self.device_info_discovery_update(config)
+        await self._subscribe_topics()
+        self.async_write_ha_state()
+
+    async def async_added_to_hass(self):
+        """Subscribe MQTT events."""
+        await super().async_added_to_hass()
+        await self._subscribe_topics()
+
+    async def async_will_remove_from_hass(self):
+        """Unsubscribe when removed."""
+        await subscription.async_unsubscribe_topics(self.hass, self._sub_state)
+        await MqttAttributes.async_will_remove_from_hass(self)
+        await MqttAvailability.async_will_remove_from_hass(self)
+
+    async def _subscribe_topics(self):
+        """(Re)Subscribe to topics."""
+        template = self._config.get(CONF_VALUE_TEMPLATE)
+        if template is not None:
+            template.hass = self.hass
+        topics = {}
+
+        @callback
+        def state_message_received(msg):
+            """Handle state MQTT message."""
+            payload = msg.payload
+            if template is not None:
+                payload = template.async_render_with_possible_json_value(
+                    payload)
+            else:
+                payload = json.loads(payload)
+            if STATE in payload and payload[STATE] in POSSIBLE_STATES:
+                self._state = POSSIBLE_STATES[payload[STATE]]
+                del payload[STATE]
+            self._state_attrs.update(payload)
+            self.async_write_ha_state()
+
+        if self._config.get(CONF_STATE_TOPIC):
+            topics['state_position_topic'] = {
+                'topic': self._config.get(CONF_STATE_TOPIC),
+                'msg_callback': state_message_received,
+                'qos': self._config[CONF_QOS]}
+        self._sub_state = await subscription.async_subscribe_topics(
+            self.hass, self._sub_state, topics)
+
+    @property
+    def name(self):
+        """Return the name of the vacuum."""
+        return self._name
+
+    @property
+    def state(self):
+        """Return state of vacuum."""
+        return self._state
+
+    @property
+    def unique_id(self):
+        """Return a unique ID."""
+        return self._unique_id
+
+    @property
+    def fan_speed(self):
+        """Return fan speed of the vacuum."""
+        if self.supported_features & SUPPORT_FAN_SPEED == 0:
+            return None
+
+        return self._state_attrs.get(FAN_SPEED, 0)
+
+    @property
+    def fan_speed_list(self):
+        """Return fan speed list of the vacuum."""
+        if self.supported_features & SUPPORT_FAN_SPEED == 0:
+            return None
+        return self._fan_speed_list
+
+    @property
+    def battery_level(self):
+        """Return battery level of the vacuum."""
+        if self.supported_features & SUPPORT_BATTERY == 0:
+            return None
+        return max(0, min(100, self._state_attrs.get(BATTERY, 0)))
+
+    @property
+    def supported_features(self):
+        """Flag supported features."""
+        return self._supported_features
+
+    async def async_start(self):
+        """Start the vacuum."""
+        if self.supported_features & SUPPORT_START == 0:
+            return None
+        mqtt.async_publish(self.hass, self._command_topic,
+                           self._config[CONF_PAYLOAD_START],
+                           self._config[CONF_QOS],
+                           self._config[CONF_RETAIN])
+
+    async def async_pause(self):
+        """Pause the vacuum."""
+        if self.supported_features & SUPPORT_PAUSE == 0:
+            return None
+        mqtt.async_publish(self.hass, self._command_topic,
+                           self._config[CONF_PAYLOAD_PAUSE],
+                           self._config[CONF_QOS],
+                           self._config[CONF_RETAIN])
+
+    async def async_stop(self, **kwargs):
+        """Stop the vacuum."""
+        if self.supported_features & SUPPORT_STOP == 0:
+            return None
+        mqtt.async_publish(self.hass, self._command_topic,
+                           self._config[CONF_PAYLOAD_STOP],
+                           self._config[CONF_QOS],
+                           self._config[CONF_RETAIN])
+
+    async def async_set_fan_speed(self, fan_speed, **kwargs):
+        """Set fan speed."""
+        if ((self.supported_features & SUPPORT_FAN_SPEED == 0) or
+                (fan_speed not in self._fan_speed_list)):
+            return None
+        mqtt.async_publish(self.hass, self._set_fan_speed_topic,
+                           fan_speed,
+                           self._config[CONF_QOS],
+                           self._config[CONF_RETAIN])
+
+    async def async_return_to_base(self, **kwargs):
+        """Tell the vacuum to return to its dock."""
+        if self.supported_features & SUPPORT_RETURN_HOME == 0:
+            return None
+        mqtt.async_publish(self.hass, self._command_topic,
+                           self._config[CONF_PAYLOAD_RETURN_TO_BASE],
+                           self._config[CONF_QOS],
+                           self._config[CONF_RETAIN])
+
+    async def async_clean_spot(self, **kwargs):
+        """Perform a spot clean-up."""
+        if self.supported_features & SUPPORT_CLEAN_SPOT == 0:
+            return None
+        mqtt.async_publish(self.hass, self._command_topic,
+                           self._config[CONF_PAYLOAD_CLEAN_SPOT],
+                           self._config[CONF_QOS],
+                           self._config[CONF_RETAIN])
+
+    async def async_locate(self, **kwargs):
+        """Locate the vacuum (usually by playing a song)."""
+        if self.supported_features & SUPPORT_LOCATE == 0:
+            return None
+        mqtt.async_publish(self.hass, self._command_topic,
+                           self._config[CONF_PAYLOAD_LOCATE],
+                           self._config[CONF_QOS],
+                           self._config[CONF_RETAIN])
+
+    async def async_send_command(self, command, params=None, **kwargs):
+        """Send a command to a vacuum cleaner."""
+        if self.supported_features & SUPPORT_SEND_COMMAND == 0:
+            return None
+        if params:
+            message = {"command": command}
+            message.update(params)
+            message = json.dumps(message)
+        else:
+            message = command
+        mqtt.async_publish(self.hass, self._send_command_topic,
+                           message,
+                           self._config[CONF_QOS],
+                           self._config[CONF_RETAIN])

tests/components/mqtt/test_legacy_vacuum.py

@@ -1,12 +1,14 @@
-"""The tests for the Mqtt vacuum platform."""
+"""The tests for the Legacy Mqtt vacuum platform."""
-import copy
+from copy import deepcopy
 import json
-import pytest
 
 from homeassistant.components import mqtt, vacuum
-from homeassistant.components.mqtt import (
+from homeassistant.components.mqtt import CONF_COMMAND_TOPIC
-    CONF_COMMAND_TOPIC, vacuum as mqttvacuum)
 from homeassistant.components.mqtt.discovery import async_start
+from homeassistant.components.mqtt.vacuum import (
+    schema_legacy as mqttvacuum, services_to_strings)
+from homeassistant.components.mqtt.vacuum.schema_legacy import (
+    ALL_SERVICES, SERVICE_TO_STRING)
 from homeassistant.components.vacuum import (
     ATTR_BATTERY_ICON, ATTR_BATTERY_LEVEL, ATTR_FAN_SPEED, ATTR_STATUS)
 from homeassistant.const import (
@@ -17,7 +19,7 @@ from tests.common import (
     MockConfigEntry, async_fire_mqtt_message, async_mock_mqtt_component)
 from tests.components.vacuum import common
 
-default_config = {
+DEFAULT_CONFIG = {
     CONF_PLATFORM: 'mqtt',
     CONF_NAME: 'mqtttest',
     CONF_COMMAND_TOPIC: 'vacuum/command',
@@ -40,115 +42,205 @@ default_config = {
 }
 
 
-@pytest.fixture
+async def test_default_supported_features(hass, mqtt_mock):
-def mock_publish(hass):
-    """Initialize components."""
-    yield hass.loop.run_until_complete(async_mock_mqtt_component(hass))
-
-
-async def test_default_supported_features(hass, mock_publish):
     """Test that the correct supported features."""
     assert await async_setup_component(hass, vacuum.DOMAIN, {
-        vacuum.DOMAIN: default_config,
+        vacuum.DOMAIN: DEFAULT_CONFIG,
     })
     entity = hass.states.get('vacuum.mqtttest')
     entity_features = \
         entity.attributes.get(mqttvacuum.CONF_SUPPORTED_FEATURES, 0)
-    assert sorted(mqttvacuum.services_to_strings(entity_features)) == \
+    assert sorted(services_to_strings(entity_features, SERVICE_TO_STRING)) == \
         sorted(['turn_on', 'turn_off', 'stop',
                 'return_home', 'battery', 'status',
                 'clean_spot'])
 
 
-async def test_all_commands(hass, mock_publish):
+async def test_all_commands(hass, mqtt_mock):
     """Test simple commands to the vacuum."""
-    default_config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+    config = deepcopy(DEFAULT_CONFIG)
-        mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES)
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        mqttvacuum.services_to_strings(ALL_SERVICES, SERVICE_TO_STRING)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
-        vacuum.DOMAIN: default_config,
+        vacuum.DOMAIN: config,
     })
 
     common.turn_on(hass, 'vacuum.mqtttest')
     await hass.async_block_till_done()
     await hass.async_block_till_done()
-    mock_publish.async_publish.assert_called_once_with(
+    mqtt_mock.async_publish.assert_called_once_with(
         'vacuum/command', 'turn_on', 0, False)
-    mock_publish.async_publish.reset_mock()
+    mqtt_mock.async_publish.reset_mock()
 
     common.turn_off(hass, 'vacuum.mqtttest')
     await hass.async_block_till_done()
     await hass.async_block_till_done()
-    mock_publish.async_publish.assert_called_once_with(
+    mqtt_mock.async_publish.assert_called_once_with(
         'vacuum/command', 'turn_off', 0, False)
-    mock_publish.async_publish.reset_mock()
+    mqtt_mock.async_publish.reset_mock()
 
     common.stop(hass, 'vacuum.mqtttest')
     await hass.async_block_till_done()
     await hass.async_block_till_done()
-    mock_publish.async_publish.assert_called_once_with(
+    mqtt_mock.async_publish.assert_called_once_with(
         'vacuum/command', 'stop', 0, False)
-    mock_publish.async_publish.reset_mock()
+    mqtt_mock.async_publish.reset_mock()
 
     common.clean_spot(hass, 'vacuum.mqtttest')
     await hass.async_block_till_done()
     await hass.async_block_till_done()
-    mock_publish.async_publish.assert_called_once_with(
+    mqtt_mock.async_publish.assert_called_once_with(
         'vacuum/command', 'clean_spot', 0, False)
-    mock_publish.async_publish.reset_mock()
+    mqtt_mock.async_publish.reset_mock()
 
     common.locate(hass, 'vacuum.mqtttest')
     await hass.async_block_till_done()
     await hass.async_block_till_done()
-    mock_publish.async_publish.assert_called_once_with(
+    mqtt_mock.async_publish.assert_called_once_with(
         'vacuum/command', 'locate', 0, False)
-    mock_publish.async_publish.reset_mock()
+    mqtt_mock.async_publish.reset_mock()
 
     common.start_pause(hass, 'vacuum.mqtttest')
     await hass.async_block_till_done()
     await hass.async_block_till_done()
-    mock_publish.async_publish.assert_called_once_with(
+    mqtt_mock.async_publish.assert_called_once_with(
         'vacuum/command', 'start_pause', 0, False)
-    mock_publish.async_publish.reset_mock()
+    mqtt_mock.async_publish.reset_mock()
 
     common.return_to_base(hass, 'vacuum.mqtttest')
     await hass.async_block_till_done()
     await hass.async_block_till_done()
-    mock_publish.async_publish.assert_called_once_with(
+    mqtt_mock.async_publish.assert_called_once_with(
         'vacuum/command', 'return_to_base', 0, False)
-    mock_publish.async_publish.reset_mock()
+    mqtt_mock.async_publish.reset_mock()
 
     common.set_fan_speed(hass, 'high', 'vacuum.mqtttest')
     await hass.async_block_till_done()
     await hass.async_block_till_done()
-    mock_publish.async_publish.assert_called_once_with(
+    mqtt_mock.async_publish.assert_called_once_with(
         'vacuum/set_fan_speed', 'high', 0, False)
-    mock_publish.async_publish.reset_mock()
+    mqtt_mock.async_publish.reset_mock()
 
     common.send_command(hass, '44 FE 93', entity_id='vacuum.mqtttest')
     await hass.async_block_till_done()
     await hass.async_block_till_done()
-    mock_publish.async_publish.assert_called_once_with(
+    mqtt_mock.async_publish.assert_called_once_with(
         'vacuum/send_command', '44 FE 93', 0, False)
-    mock_publish.async_publish.reset_mock()
+    mqtt_mock.async_publish.reset_mock()
 
     common.send_command(hass, '44 FE 93', {"key": "value"},
                         entity_id='vacuum.mqtttest')
     await hass.async_block_till_done()
     await hass.async_block_till_done()
-    assert json.loads(mock_publish.async_publish.mock_calls[-1][1][1]) == {
+    assert json.loads(mqtt_mock.async_publish.mock_calls[-1][1][1]) == {
+        "command": "44 FE 93",
+        "key": "value"
+    }
+
+    common.send_command(hass, '44 FE 93', {"key": "value"},
+                        entity_id='vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    assert json.loads(mqtt_mock.async_publish.mock_calls[-1][1][1]) == {
         "command": "44 FE 93",
         "key": "value"
     }
 
 
-async def test_status(hass, mock_publish):
+async def test_commands_without_supported_features(hass, mqtt_mock):
-    """Test status updates from the vacuum."""
+    """Test commands which are not supported by the vacuum."""
-    default_config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+    config = deepcopy(DEFAULT_CONFIG)
-        mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES)
+    services = mqttvacuum.STRING_TO_SERVICE["status"]
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        services_to_strings(
+            services, SERVICE_TO_STRING)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
-        vacuum.DOMAIN: default_config,
+        vacuum.DOMAIN: config,
+    })
+
+    common.turn_on(hass, 'vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    common.turn_off(hass, 'vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    common.stop(hass, 'vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    common.clean_spot(hass, 'vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    common.locate(hass, 'vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    common.start_pause(hass, 'vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    common.return_to_base(hass, 'vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    common.set_fan_speed(hass, 'high', 'vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    common.send_command(hass, '44 FE 93', entity_id='vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+
+async def test_attributes_without_supported_features(hass, mqtt_mock):
+    """Test attributes which are not supported by the vacuum."""
+    config = deepcopy(DEFAULT_CONFIG)
+    services = mqttvacuum.STRING_TO_SERVICE["turn_on"]
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        services_to_strings(
+            services, SERVICE_TO_STRING)
+
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: config,
+    })
+
+    state = hass.states.get('vacuum.mqtttest')
+    assert STATE_OFF == state.state
+    assert state.attributes.get(ATTR_BATTERY_LEVEL) is None
+    assert state.attributes.get(ATTR_BATTERY_ICON) is None
+
+
+async def test_status(hass, mqtt_mock):
+    """Test status updates from the vacuum."""
+    config = deepcopy(DEFAULT_CONFIG)
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        mqttvacuum.services_to_strings(ALL_SERVICES, SERVICE_TO_STRING)
+
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: config,
     })
 
     message = """{
@@ -162,11 +254,10 @@ async def test_status(hass, mock_publish):
     await hass.async_block_till_done()
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.mqtttest')
-    assert STATE_ON == state.state
+    assert state.state == STATE_ON
-    assert 'mdi:battery-50' == \
+    assert state.attributes.get(ATTR_BATTERY_ICON) == 'mdi:battery-50'
-        state.attributes.get(ATTR_BATTERY_ICON)
+    assert state.attributes.get(ATTR_BATTERY_LEVEL) == 54
-    assert 54 == state.attributes.get(ATTR_BATTERY_LEVEL)
+    assert state.attributes.get(ATTR_FAN_SPEED) == 'max'
-    assert 'max' == state.attributes.get(ATTR_FAN_SPEED)
 
     message = """{
         "battery_level": 61,
@@ -180,20 +271,20 @@ async def test_status(hass, mock_publish):
     await hass.async_block_till_done()
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.mqtttest')
-    assert STATE_OFF == state.state
+    assert state.state == STATE_OFF
-    assert 'mdi:battery-charging-60' == \
+    assert state.attributes.get(ATTR_BATTERY_ICON) == 'mdi:battery-charging-60'
-        state.attributes.get(ATTR_BATTERY_ICON)
+    assert state.attributes.get(ATTR_BATTERY_LEVEL) == 61
-    assert 61 == state.attributes.get(ATTR_BATTERY_LEVEL)
+    assert state.attributes.get(ATTR_FAN_SPEED) == 'min'
-    assert 'min' == state.attributes.get(ATTR_FAN_SPEED)
 
 
-async def test_status_battery(hass, mock_publish):
+async def test_status_battery(hass, mqtt_mock):
     """Test status updates from the vacuum."""
-    default_config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+    config = deepcopy(DEFAULT_CONFIG)
-        mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES)
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        mqttvacuum.services_to_strings(ALL_SERVICES, SERVICE_TO_STRING)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
-        vacuum.DOMAIN: default_config,
+        vacuum.DOMAIN: config,
     })
 
     message = """{
@@ -203,17 +294,17 @@ async def test_status_battery(hass, mock_publish):
     await hass.async_block_till_done()
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.mqtttest')
-    assert 'mdi:battery-50' == \
+    assert state.attributes.get(ATTR_BATTERY_ICON) == 'mdi:battery-50'
-        state.attributes.get(ATTR_BATTERY_ICON)
 
 
-async def test_status_cleaning(hass, mock_publish):
+async def test_status_cleaning(hass, mqtt_mock):
     """Test status updates from the vacuum."""
-    default_config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+    config = deepcopy(DEFAULT_CONFIG)
-        mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES)
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        mqttvacuum.services_to_strings(ALL_SERVICES, SERVICE_TO_STRING)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
-        vacuum.DOMAIN: default_config,
+        vacuum.DOMAIN: config,
     })
 
     message = """{
@@ -223,16 +314,17 @@ async def test_status_cleaning(hass, mock_publish):
     await hass.async_block_till_done()
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.mqtttest')
-    assert STATE_ON == state.state
+    assert state.state == STATE_ON
 
 
-async def test_status_docked(hass, mock_publish):
+async def test_status_docked(hass, mqtt_mock):
     """Test status updates from the vacuum."""
-    default_config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+    config = deepcopy(DEFAULT_CONFIG)
-        mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES)
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        mqttvacuum.services_to_strings(ALL_SERVICES, SERVICE_TO_STRING)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
-        vacuum.DOMAIN: default_config,
+        vacuum.DOMAIN: config,
     })
 
     message = """{
@@ -242,16 +334,17 @@ async def test_status_docked(hass, mock_publish):
     await hass.async_block_till_done()
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.mqtttest')
-    assert STATE_OFF == state.state
+    assert state.state == STATE_OFF
 
 
-async def test_status_charging(hass, mock_publish):
+async def test_status_charging(hass, mqtt_mock):
     """Test status updates from the vacuum."""
-    default_config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+    config = deepcopy(DEFAULT_CONFIG)
-        mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES)
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        mqttvacuum.services_to_strings(ALL_SERVICES, SERVICE_TO_STRING)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
-        vacuum.DOMAIN: default_config,
+        vacuum.DOMAIN: config,
     })
 
     message = """{
@@ -261,17 +354,17 @@ async def test_status_charging(hass, mock_publish):
     await hass.async_block_till_done()
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.mqtttest')
-    assert 'mdi:battery-outline' == \
+    assert state.attributes.get(ATTR_BATTERY_ICON) == 'mdi:battery-outline'
-        state.attributes.get(ATTR_BATTERY_ICON)
 
 
-async def test_status_fan_speed(hass, mock_publish):
+async def test_status_fan_speed(hass, mqtt_mock):
     """Test status updates from the vacuum."""
-    default_config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+    config = deepcopy(DEFAULT_CONFIG)
-        mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES)
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        mqttvacuum.services_to_strings(ALL_SERVICES, SERVICE_TO_STRING)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
-        vacuum.DOMAIN: default_config,
+        vacuum.DOMAIN: config,
     })
 
     message = """{
@@ -281,16 +374,17 @@ async def test_status_fan_speed(hass, mock_publish):
     await hass.async_block_till_done()
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.mqtttest')
-    assert 'max' == state.attributes.get(ATTR_FAN_SPEED)
+    assert state.attributes.get(ATTR_FAN_SPEED) == 'max'
 
 
-async def test_status_error(hass, mock_publish):
+async def test_status_error(hass, mqtt_mock):
     """Test status updates from the vacuum."""
-    default_config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+    config = deepcopy(DEFAULT_CONFIG)
-        mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES)
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        mqttvacuum.services_to_strings(ALL_SERVICES, SERVICE_TO_STRING)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
-        vacuum.DOMAIN: default_config,
+        vacuum.DOMAIN: config,
     })
 
     message = """{
@@ -299,7 +393,7 @@ async def test_status_error(hass, mock_publish):
     async_fire_mqtt_message(hass, 'vacuum/state', message)
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.mqtttest')
-    assert 'Error: Error1' == state.attributes.get(ATTR_STATUS)
+    assert state.attributes.get(ATTR_STATUS) == 'Error: Error1'
 
     message = """{
         "error": ""
@@ -307,49 +401,50 @@ async def test_status_error(hass, mock_publish):
     async_fire_mqtt_message(hass, 'vacuum/state', message)
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.mqtttest')
-    assert 'Stopped' == state.attributes.get(ATTR_STATUS)
+    assert state.attributes.get(ATTR_STATUS) == 'Stopped'
 
 
-async def test_battery_template(hass, mock_publish):
+async def test_battery_template(hass, mqtt_mock):
     """Test that you can use non-default templates for battery_level."""
-    default_config.update({
+    config = deepcopy(DEFAULT_CONFIG)
+    config.update({
         mqttvacuum.CONF_SUPPORTED_FEATURES:
-            mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES),
+            mqttvacuum.services_to_strings(ALL_SERVICES, SERVICE_TO_STRING),
         mqttvacuum.CONF_BATTERY_LEVEL_TOPIC: "retroroomba/battery_level",
         mqttvacuum.CONF_BATTERY_LEVEL_TEMPLATE: "{{ value }}"
     })
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
-        vacuum.DOMAIN: default_config,
+        vacuum.DOMAIN: config,
     })
 
     async_fire_mqtt_message(hass, 'retroroomba/battery_level', '54')
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.mqtttest')
-    assert 54 == state.attributes.get(ATTR_BATTERY_LEVEL)
+    assert state.attributes.get(ATTR_BATTERY_LEVEL) == 54
-    assert state.attributes.get(ATTR_BATTERY_ICON) == \
+    assert state.attributes.get(ATTR_BATTERY_ICON) == 'mdi:battery-50'
-        'mdi:battery-50'
 
 
-async def test_status_invalid_json(hass, mock_publish):
+async def test_status_invalid_json(hass, mqtt_mock):
     """Test to make sure nothing breaks if the vacuum sends bad JSON."""
-    default_config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+    config = deepcopy(DEFAULT_CONFIG)
-        mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES)
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        mqttvacuum.services_to_strings(ALL_SERVICES, SERVICE_TO_STRING)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
-        vacuum.DOMAIN: default_config,
+        vacuum.DOMAIN: config,
     })
 
     async_fire_mqtt_message(hass, 'vacuum/state', '{"asdfasas false}')
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.mqtttest')
-    assert STATE_OFF == state.state
+    assert state.state == STATE_OFF
-    assert "Stopped" == state.attributes.get(ATTR_STATUS)
+    assert state.attributes.get(ATTR_STATUS) == "Stopped"
 
 
-async def test_missing_battery_template(hass, mock_publish):
+async def test_missing_battery_template(hass, mqtt_mock):
     """Test to make sure missing template is not allowed."""
-    config = copy.deepcopy(default_config)
+    config = deepcopy(DEFAULT_CONFIG)
     config.pop(mqttvacuum.CONF_BATTERY_LEVEL_TEMPLATE)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
@@ -360,9 +455,9 @@ async def test_missing_battery_template(hass, mock_publish):
     assert state is None
 
 
-async def test_missing_charging_template(hass, mock_publish):
+async def test_missing_charging_template(hass, mqtt_mock):
     """Test to make sure missing template is not allowed."""
-    config = copy.deepcopy(default_config)
+    config = deepcopy(DEFAULT_CONFIG)
     config.pop(mqttvacuum.CONF_CHARGING_TEMPLATE)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
@@ -373,9 +468,9 @@ async def test_missing_charging_template(hass, mock_publish):
     assert state is None
 
 
-async def test_missing_cleaning_template(hass, mock_publish):
+async def test_missing_cleaning_template(hass, mqtt_mock):
     """Test to make sure missing template is not allowed."""
-    config = copy.deepcopy(default_config)
+    config = deepcopy(DEFAULT_CONFIG)
     config.pop(mqttvacuum.CONF_CLEANING_TEMPLATE)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
@@ -386,9 +481,9 @@ async def test_missing_cleaning_template(hass, mock_publish):
     assert state is None
 
 
-async def test_missing_docked_template(hass, mock_publish):
+async def test_missing_docked_template(hass, mqtt_mock):
     """Test to make sure missing template is not allowed."""
-    config = copy.deepcopy(default_config)
+    config = deepcopy(DEFAULT_CONFIG)
     config.pop(mqttvacuum.CONF_DOCKED_TEMPLATE)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
@@ -399,9 +494,9 @@ async def test_missing_docked_template(hass, mock_publish):
     assert state is None
 
 
-async def test_missing_error_template(hass, mock_publish):
+async def test_missing_error_template(hass, mqtt_mock):
     """Test to make sure missing template is not allowed."""
-    config = copy.deepcopy(default_config)
+    config = deepcopy(DEFAULT_CONFIG)
     config.pop(mqttvacuum.CONF_ERROR_TEMPLATE)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
@@ -412,9 +507,9 @@ async def test_missing_error_template(hass, mock_publish):
     assert state is None
 
 
-async def test_missing_fan_speed_template(hass, mock_publish):
+async def test_missing_fan_speed_template(hass, mqtt_mock):
     """Test to make sure missing template is not allowed."""
-    config = copy.deepcopy(default_config)
+    config = deepcopy(DEFAULT_CONFIG)
     config.pop(mqttvacuum.CONF_FAN_SPEED_TEMPLATE)
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
@@ -425,14 +520,15 @@ async def test_missing_fan_speed_template(hass, mock_publish):
     assert state is None
 
 
-async def test_default_availability_payload(hass, mock_publish):
+async def test_default_availability_payload(hass, mqtt_mock):
     """Test availability by default payload with defined topic."""
-    default_config.update({
+    config = deepcopy(DEFAULT_CONFIG)
+    config.update({
         'availability_topic': 'availability-topic'
     })
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
-        vacuum.DOMAIN: default_config,
+        vacuum.DOMAIN: config,
     })
 
     state = hass.states.get('vacuum.mqtttest')
@@ -453,16 +549,17 @@ async def test_default_availability_payload(hass, mock_publish):
     assert STATE_UNAVAILABLE == state.state
 
 
-async def test_custom_availability_payload(hass, mock_publish):
+async def test_custom_availability_payload(hass, mqtt_mock):
     """Test availability by custom payload with defined topic."""
-    default_config.update({
+    config = deepcopy(DEFAULT_CONFIG)
+    config.update({
         'availability_topic': 'availability-topic',
         'payload_available': 'good',
         'payload_not_available': 'nogood'
     })
 
     assert await async_setup_component(hass, vacuum.DOMAIN, {
-        vacuum.DOMAIN: default_config,
+        vacuum.DOMAIN: config,
     })
 
     state = hass.states.get('vacuum.mqtttest')
@@ -483,7 +580,7 @@ async def test_custom_availability_payload(hass, mock_publish):
     assert STATE_UNAVAILABLE == state.state
 
 
-async def test_discovery_removal_vacuum(hass, mock_publish):
+async def test_discovery_removal_vacuum(hass, mqtt_mock):
     """Test removal of discovered vacuum."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)
     await async_start(hass, 'homeassistant', {}, entry)
@@ -543,7 +640,7 @@ async def test_discovery_broken(hass, mqtt_mock, caplog):
     assert state is None
 
 
-async def test_discovery_update_vacuum(hass, mock_publish):
+async def test_discovery_update_vacuum(hass, mqtt_mock):
     """Test update of discovered vacuum."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)
     await async_start(hass, 'homeassistant', {}, entry)
@@ -592,7 +689,7 @@ async def test_setting_attribute_via_mqtt_json_message(hass, mqtt_mock):
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.test')
 
-    assert '100' == state.attributes.get('val')
+    assert state.attributes.get('val') == '100'
 
 
 async def test_update_with_json_attrs_not_dict(hass, mqtt_mock, caplog):
@@ -614,7 +711,7 @@ async def test_update_with_json_attrs_not_dict(hass, mqtt_mock, caplog):
     assert 'JSON result was not a dictionary' in caplog.text
 
 
-async def test_update_with_json_attrs_bad_JSON(hass, mqtt_mock, caplog):
+async def test_update_with_json_attrs_bad_json(hass, mqtt_mock, caplog):
     """Test attributes get extracted from a JSON result."""
     assert await async_setup_component(hass, vacuum.DOMAIN, {
         vacuum.DOMAIN: {
@@ -654,7 +751,7 @@ async def test_discovery_update_attr(hass, mqtt_mock, caplog):
     await hass.async_block_till_done()
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.beer')
-    assert '100' == state.attributes.get('val')
+    assert state.attributes.get('val') == '100'
 
     # Change json_attributes_topic
     async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
@@ -667,17 +764,17 @@ async def test_discovery_update_attr(hass, mqtt_mock, caplog):
     await hass.async_block_till_done()
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.beer')
-    assert '100' == state.attributes.get('val')
+    assert state.attributes.get('val') == '100'
 
     # Verify we are subscribing to the new topic
     async_fire_mqtt_message(hass, 'attr-topic2', '{ "val": "75" }')
     await hass.async_block_till_done()
     await hass.async_block_till_done()
     state = hass.states.get('vacuum.beer')
-    assert '75' == state.attributes.get('val')
+    assert state.attributes.get('val') == '75'
 
 
-async def test_unique_id(hass, mock_publish):
+async def test_unique_id(hass, mqtt_mock):
     """Test unique id option only creates one vacuum per unique_id."""
     await async_mock_mqtt_component(hass)
     assert await async_setup_component(hass, vacuum.DOMAIN, {
@@ -702,7 +799,7 @@ async def test_unique_id(hass, mock_publish):
     # all vacuums group is 1, unique id created is 1
 
 
-async def test_entity_device_info_with_identifier(hass, mock_publish):
+async def test_entity_device_info_with_identifier(hass, mqtt_mock):
     """Test MQTT vacuum device registry integration."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)
     entry.add_to_hass(hass)

tests/components/mqtt/test_state_vacuum.py

@@ -0,0 +1,685 @@
+"""The tests for the State vacuum Mqtt platform."""
+from copy import deepcopy
+import json
+
+from homeassistant.components import mqtt, vacuum
+from homeassistant.components.mqtt import CONF_COMMAND_TOPIC, CONF_STATE_TOPIC
+from homeassistant.components.mqtt.discovery import async_start
+from homeassistant.components.mqtt.vacuum import (
+    CONF_SCHEMA, schema_state as mqttvacuum, services_to_strings)
+from homeassistant.components.mqtt.vacuum.schema_state import SERVICE_TO_STRING
+from homeassistant.components.vacuum import (
+    ATTR_BATTERY_ICON, ATTR_BATTERY_LEVEL, ATTR_FAN_SPEED, ATTR_FAN_SPEED_LIST,
+    DOMAIN, SERVICE_CLEAN_SPOT, SERVICE_LOCATE, SERVICE_PAUSE,
+    SERVICE_RETURN_TO_BASE, SERVICE_START, SERVICE_STOP, STATE_CLEANING,
+    STATE_DOCKED)
+from homeassistant.const import (
+    CONF_NAME, CONF_PLATFORM, STATE_UNAVAILABLE, STATE_UNKNOWN)
+from homeassistant.setup import async_setup_component
+
+from tests.common import (
+    MockConfigEntry, async_fire_mqtt_message, async_mock_mqtt_component)
+from tests.components.vacuum import common
+
+COMMAND_TOPIC = 'vacuum/command'
+SEND_COMMAND_TOPIC = 'vacuum/send_command'
+STATE_TOPIC = 'vacuum/state'
+
+DEFAULT_CONFIG = {
+    CONF_PLATFORM: 'mqtt',
+    CONF_SCHEMA: 'state',
+    CONF_NAME: 'mqtttest',
+    CONF_COMMAND_TOPIC: COMMAND_TOPIC,
+    mqttvacuum.CONF_SEND_COMMAND_TOPIC: SEND_COMMAND_TOPIC,
+    CONF_STATE_TOPIC: STATE_TOPIC,
+    mqttvacuum.CONF_SET_FAN_SPEED_TOPIC: 'vacuum/set_fan_speed',
+    mqttvacuum.CONF_FAN_SPEED_LIST: ['min', 'medium', 'high', 'max'],
+}
+
+
+async def test_default_supported_features(hass, mqtt_mock):
+    """Test that the correct supported features."""
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: DEFAULT_CONFIG,
+    })
+    entity = hass.states.get('vacuum.mqtttest')
+    entity_features = \
+        entity.attributes.get(mqttvacuum.CONF_SUPPORTED_FEATURES, 0)
+    assert sorted(services_to_strings(entity_features, SERVICE_TO_STRING)) == \
+        sorted(['start', 'stop',
+                'return_home', 'battery', 'status',
+                'clean_spot'])
+
+
+async def test_all_commands(hass, mqtt_mock):
+    """Test simple commands send to the vacuum."""
+    config = deepcopy(DEFAULT_CONFIG)
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        services_to_strings(
+            mqttvacuum.ALL_SERVICES, SERVICE_TO_STRING)
+
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: config,
+    })
+
+    await hass.services.async_call(
+        DOMAIN, SERVICE_START, blocking=True)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_called_once_with(
+        COMMAND_TOPIC, 'start', 0, False)
+    mqtt_mock.async_publish.reset_mock()
+
+    await hass.services.async_call(
+        DOMAIN, SERVICE_STOP, blocking=True)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_called_once_with(
+        COMMAND_TOPIC, 'stop', 0, False)
+    mqtt_mock.async_publish.reset_mock()
+
+    await hass.services.async_call(
+        DOMAIN, SERVICE_PAUSE, blocking=True)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_called_once_with(
+        COMMAND_TOPIC, 'pause', 0, False)
+    mqtt_mock.async_publish.reset_mock()
+
+    await hass.services.async_call(
+        DOMAIN, SERVICE_LOCATE, blocking=True)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_called_once_with(
+        COMMAND_TOPIC, 'locate', 0, False)
+    mqtt_mock.async_publish.reset_mock()
+
+    await hass.services.async_call(
+        DOMAIN, SERVICE_CLEAN_SPOT, blocking=True)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_called_once_with(
+        COMMAND_TOPIC, 'clean_spot', 0, False)
+    mqtt_mock.async_publish.reset_mock()
+
+    await hass.services.async_call(
+        DOMAIN, SERVICE_RETURN_TO_BASE, blocking=True)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_called_once_with(
+        COMMAND_TOPIC, 'return_to_base', 0, False)
+    mqtt_mock.async_publish.reset_mock()
+
+    common.set_fan_speed(hass, 'medium', 'vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_called_once_with(
+        'vacuum/set_fan_speed', 'medium', 0, False)
+    mqtt_mock.async_publish.reset_mock()
+
+    common.send_command(hass, '44 FE 93', entity_id='vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_called_once_with(
+        'vacuum/send_command', '44 FE 93', 0, False)
+    mqtt_mock.async_publish.reset_mock()
+
+    common.send_command(hass, '44 FE 93', {"key": "value"},
+                        entity_id='vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    assert json.loads(mqtt_mock.async_publish.mock_calls[-1][1][1]) == {
+        "command": "44 FE 93",
+        "key": "value"
+    }
+
+
+async def test_commands_without_supported_features(hass, mqtt_mock):
+    """Test commands which are not supported by the vacuum."""
+    config = deepcopy(DEFAULT_CONFIG)
+    services = mqttvacuum.STRING_TO_SERVICE["status"]
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        services_to_strings(
+            services, SERVICE_TO_STRING)
+
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: config,
+    })
+
+    await hass.services.async_call(
+        DOMAIN, SERVICE_START, blocking=True)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    await hass.services.async_call(
+        DOMAIN, SERVICE_PAUSE, blocking=True)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    await hass.services.async_call(
+        DOMAIN, SERVICE_STOP, blocking=True)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    await hass.services.async_call(
+        DOMAIN, SERVICE_RETURN_TO_BASE, blocking=True)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    await hass.services.async_call(
+        DOMAIN, SERVICE_LOCATE, blocking=True)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    await hass.services.async_call(
+        DOMAIN, SERVICE_CLEAN_SPOT, blocking=True)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    common.set_fan_speed(hass, 'medium', 'vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+    mqtt_mock.async_publish.reset_mock()
+
+    common.send_command(hass, '44 FE 93', {"key": "value"},
+                        entity_id='vacuum.mqtttest')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    mqtt_mock.async_publish.assert_not_called()
+
+
+async def test_status(hass, mqtt_mock):
+    """Test status updates from the vacuum."""
+    config = deepcopy(DEFAULT_CONFIG)
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        services_to_strings(mqttvacuum.ALL_SERVICES, SERVICE_TO_STRING)
+
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: config,
+    })
+
+    message = """{
+        "battery_level": 54,
+        "state": "cleaning",
+        "fan_speed": "max"
+    }"""
+    async_fire_mqtt_message(hass, 'vacuum/state', message)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.mqtttest')
+    assert state.state == STATE_CLEANING
+    assert state.attributes.get(ATTR_BATTERY_LEVEL) == 54
+    assert state.attributes.get(ATTR_BATTERY_ICON) == 'mdi:battery-50'
+    assert state.attributes.get(ATTR_FAN_SPEED) == 'max'
+
+    message = """{
+        "battery_level": 61,
+        "state": "docked",
+        "fan_speed": "min"
+    }"""
+
+    async_fire_mqtt_message(hass, 'vacuum/state', message)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.mqtttest')
+    assert state.state == STATE_DOCKED
+    assert state.attributes.get(ATTR_BATTERY_ICON) == 'mdi:battery-charging-60'
+    assert state.attributes.get(ATTR_BATTERY_LEVEL) == 61
+    assert state.attributes.get(ATTR_FAN_SPEED) == 'min'
+    assert state.attributes.get(ATTR_FAN_SPEED_LIST) == ['min', 'medium',
+                                                         'high', 'max']
+
+
+async def test_no_fan_vacuum(hass, mqtt_mock):
+    """Test status updates from the vacuum when fan is not supported."""
+    config = deepcopy(DEFAULT_CONFIG)
+    del config[mqttvacuum.CONF_FAN_SPEED_LIST]
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        services_to_strings(mqttvacuum.DEFAULT_SERVICES, SERVICE_TO_STRING)
+
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: config,
+    })
+
+    message = """{
+        "battery_level": 54,
+        "state": "cleaning"
+    }"""
+    async_fire_mqtt_message(hass, 'vacuum/state', message)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.mqtttest')
+    assert state.state == STATE_CLEANING
+    assert state.attributes.get(ATTR_FAN_SPEED) is None
+    assert state.attributes.get(ATTR_FAN_SPEED_LIST) is None
+    assert state.attributes.get(ATTR_BATTERY_LEVEL) == 54
+    assert state.attributes.get(ATTR_BATTERY_ICON) == 'mdi:battery-50'
+
+    message = """{
+        "battery_level": 54,
+        "state": "cleaning",
+        "fan_speed": "max"
+    }"""
+    async_fire_mqtt_message(hass, 'vacuum/state', message)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.mqtttest')
+
+    assert state.state == STATE_CLEANING
+    assert state.attributes.get(ATTR_FAN_SPEED) is None
+    assert state.attributes.get(ATTR_FAN_SPEED_LIST) is None
+
+    assert state.attributes.get(ATTR_BATTERY_LEVEL) == 54
+    assert state.attributes.get(ATTR_BATTERY_ICON) == 'mdi:battery-50'
+
+    message = """{
+        "battery_level": 61,
+        "state": "docked"
+    }"""
+
+    async_fire_mqtt_message(hass, 'vacuum/state', message)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.mqtttest')
+    assert state.state == STATE_DOCKED
+    assert state.attributes.get(ATTR_BATTERY_ICON) == 'mdi:battery-charging-60'
+    assert state.attributes.get(ATTR_BATTERY_LEVEL) == 61
+
+
+async def test_status_invalid_json(hass, mqtt_mock):
+    """Test to make sure nothing breaks if the vacuum sends bad JSON."""
+    config = deepcopy(DEFAULT_CONFIG)
+    config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        mqttvacuum.services_to_strings(
+            mqttvacuum.ALL_SERVICES, SERVICE_TO_STRING)
+
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: config,
+    })
+
+    async_fire_mqtt_message(hass, 'vacuum/state', '{"asdfasas false}')
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.mqtttest')
+    assert state.state == STATE_UNKNOWN
+
+
+async def test_default_availability_payload(hass, mqtt_mock):
+    """Test availability by default payload with defined topic."""
+    config = deepcopy(DEFAULT_CONFIG)
+    config.update({
+        'availability_topic': 'availability-topic'
+    })
+
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: config,
+    })
+
+    state = hass.states.get('vacuum.mqtttest')
+    assert state.state == STATE_UNAVAILABLE
+
+    async_fire_mqtt_message(hass, 'availability-topic', 'online')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('vacuum.mqtttest')
+    assert STATE_UNAVAILABLE != state.state
+
+    async_fire_mqtt_message(hass, 'availability-topic', 'offline')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('vacuum.mqtttest')
+    assert state.state == STATE_UNAVAILABLE
+
+
+async def test_custom_availability_payload(hass, mqtt_mock):
+    """Test availability by custom payload with defined topic."""
+    config = deepcopy(DEFAULT_CONFIG)
+    config.update({
+        'availability_topic': 'availability-topic',
+        'payload_available': 'good',
+        'payload_not_available': 'nogood'
+    })
+
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: config,
+    })
+
+    state = hass.states.get('vacuum.mqtttest')
+    assert state.state == STATE_UNAVAILABLE
+
+    async_fire_mqtt_message(hass, 'availability-topic', 'good')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('vacuum.mqtttest')
+    assert state.state != STATE_UNAVAILABLE
+
+    async_fire_mqtt_message(hass, 'availability-topic', 'nogood')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('vacuum.mqtttest')
+    assert state.state == STATE_UNAVAILABLE
+
+
+async def test_discovery_removal_vacuum(hass, mqtt_mock):
+    """Test removal of discovered vacuum."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    await async_start(hass, 'homeassistant', {}, entry)
+
+    data = (
+        '{ "name": "Beer",'
+        '  "command_topic": "test_topic",'
+        '  "component": "state" }'
+    )
+
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
+                            data)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('vacuum.beer')
+    assert state is not None
+    assert state.name == 'Beer'
+
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config', '')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('vacuum.beer')
+    assert state is None
+
+
+async def test_discovery_broken(hass, mqtt_mock, caplog):
+    """Test handling of bad discovery message."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    await async_start(hass, 'homeassistant', {}, entry)
+
+    data1 = (
+        '{ "name": "Beer",'
+        '  "command_topic": "test_topic#",'
+        ' "component": "state" }'
+    )
+    data2 = (
+        '{ "name": "Milk",'
+        '  "command_topic": "test_topic",'
+        ' "component": "state" }'
+    )
+
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
+                            data1)
+    await hass.async_block_till_done()
+
+    state = hass.states.get('vacuum.beer')
+    assert state is None
+
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
+                            data2)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('vacuum.milk')
+    assert state is not None
+    assert state.name == 'Milk'
+    state = hass.states.get('vacuum.beer')
+    assert state is None
+
+
+async def test_discovery_update_vacuum(hass, mqtt_mock):
+    """Test update of discovered vacuum."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    await async_start(hass, 'homeassistant', {}, entry)
+
+    data1 = (
+        '{ "name": "Beer",'
+        '  "command_topic": "test_topic",'
+        '"component": "state" }'
+    )
+    data2 = (
+        '{ "name": "Milk",'
+        '  "command_topic": "test_topic",'
+        '  "component": "state"}'
+    )
+
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
+                            data1)
+    await hass.async_block_till_done()
+
+    state = hass.states.get('vacuum.beer')
+    assert state is not None
+    assert state.name == 'Beer'
+
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
+                            data2)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('vacuum.beer')
+    assert state is not None
+    assert state.name == 'Milk'
+    state = hass.states.get('vacuum.milk')
+    assert state is None
+
+
+async def test_setting_attribute_via_mqtt_json_message(hass, mqtt_mock):
+    """Test the setting of attribute via MQTT with JSON payload."""
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: {
+            'platform': 'mqtt',
+            'name': 'test',
+            'state_topic': 'test-topic',
+            'json_attributes_topic': 'attr-topic'
+        }
+    })
+
+    async_fire_mqtt_message(hass, 'attr-topic', '{ "val": "100" }')
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.test')
+
+    assert state.attributes.get('val') == '100'
+
+
+async def test_update_with_json_attrs_not_dict(hass, mqtt_mock, caplog):
+    """Test attributes get extracted from a JSON result."""
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: {
+            'platform': 'mqtt',
+            'name': 'test',
+            'state_topic': 'test-topic',
+            'json_attributes_topic': 'attr-topic'
+        }
+    })
+
+    async_fire_mqtt_message(hass, 'attr-topic', '[ "list", "of", "things"]')
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.test')
+
+    assert state.attributes.get('val') is None
+    assert 'JSON result was not a dictionary' in caplog.text
+
+
+async def test_update_with_json_attrs_bad_json(hass, mqtt_mock, caplog):
+    """Test attributes get extracted from a JSON result."""
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: {
+            'platform': 'mqtt',
+            'name': 'test',
+            'state_topic': 'test-topic',
+            'json_attributes_topic': 'attr-topic'
+        }
+    })
+
+    async_fire_mqtt_message(hass, 'attr-topic', 'This is not JSON')
+    await hass.async_block_till_done()
+
+    state = hass.states.get('vacuum.test')
+    assert state.attributes.get('val') is None
+    assert 'Erroneous JSON: This is not JSON' in caplog.text
+
+
+async def test_discovery_update_attr(hass, mqtt_mock, caplog):
+    """Test update of discovered MQTTAttributes."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    await async_start(hass, 'homeassistant', {}, entry)
+    data1 = (
+        '{ "name": "Beer",'
+        '  "command_topic": "test_topic",'
+        '  "json_attributes_topic": "attr-topic1" }'
+    )
+    data2 = (
+        '{ "name": "Beer",'
+        '  "command_topic": "test_topic",'
+        '  "json_attributes_topic": "attr-topic2" }'
+    )
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
+                            data1)
+    await hass.async_block_till_done()
+    async_fire_mqtt_message(hass, 'attr-topic1', '{ "val": "100" }')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.beer')
+    assert state.attributes.get('val') == '100'
+
+    # Change json_attributes_topic
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
+                            data2)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    # Verify we are no longer subscribing to the old topic
+    async_fire_mqtt_message(hass, 'attr-topic1', '{ "val": "50" }')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.beer')
+    assert state.attributes.get('val') == '100'
+
+    # Verify we are subscribing to the new topic
+    async_fire_mqtt_message(hass, 'attr-topic2', '{ "val": "75" }')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.beer')
+    assert state.attributes.get('val') == '75'
+
+
+async def test_unique_id(hass, mqtt_mock):
+    """Test unique id option only creates one vacuum per unique_id."""
+    await async_mock_mqtt_component(hass)
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: [{
+            'platform': 'mqtt',
+            'name': 'Test 1',
+            'command_topic': 'command-topic',
+            'unique_id': 'TOTALLY_UNIQUE'
+        }, {
+            'platform': 'mqtt',
+            'name': 'Test 2',
+            'command_topic': 'command-topic',
+            'unique_id': 'TOTALLY_UNIQUE'
+        }]
+    })
+
+    async_fire_mqtt_message(hass, 'test-topic', 'payload')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids()) == 2
+    # all vacuums group is 1, unique id created is 1
+
+
+async def test_entity_device_info_with_identifier(hass, mqtt_mock):
+    """Test MQTT vacuum device registry integration."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    entry.add_to_hass(hass)
+    await async_start(hass, 'homeassistant', {}, entry)
+    registry = await hass.helpers.device_registry.async_get_registry()
+
+    data = json.dumps({
+        'platform': 'mqtt',
+        'name': 'Test 1',
+        'command_topic': 'test-command-topic',
+        'device': {
+            'identifiers': ['helloworld'],
+            'connections': [
+                ["mac", "02:5b:26:a8:dc:12"],
+            ],
+            'manufacturer': 'Whatever',
+            'name': 'Beer',
+            'model': 'Glass',
+            'sw_version': '0.1-beta',
+        },
+        'unique_id': 'veryunique'
+    })
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
+                            data)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    device = registry.async_get_device({('mqtt', 'helloworld')}, set())
+    assert device is not None
+    assert device.identifiers == {('mqtt', 'helloworld')}
+    assert device.connections == {('mac', "02:5b:26:a8:dc:12")}
+    assert device.manufacturer == 'Whatever'
+    assert device.name == 'Beer'
+    assert device.model == 'Glass'
+    assert device.sw_version == '0.1-beta'
+
+
+async def test_entity_device_info_update(hass, mqtt_mock):
+    """Test device registry update."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    entry.add_to_hass(hass)
+    await async_start(hass, 'homeassistant', {}, entry)
+    registry = await hass.helpers.device_registry.async_get_registry()
+
+    config = {
+        'platform': 'mqtt',
+        'name': 'Test 1',
+        'state_topic': 'test-topic',
+        'command_topic': 'test-command-topic',
+        'device': {
+            'identifiers': ['helloworld'],
+            'connections': [
+                ["mac", "02:5b:26:a8:dc:12"],
+            ],
+            'manufacturer': 'Whatever',
+            'name': 'Beer',
+            'model': 'Glass',
+            'sw_version': '0.1-beta',
+        },
+        'unique_id': 'veryunique'
+    }
+
+    data = json.dumps(config)
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
+                            data)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    device = registry.async_get_device({('mqtt', 'helloworld')}, set())
+    assert device is not None
+    assert device.name == 'Beer'
+
+    config['device']['name'] = 'Milk'
+    data = json.dumps(config)
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
+                            data)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    device = registry.async_get_device({('mqtt', 'helloworld')}, set())
+    assert device is not None
+    assert device.name == 'Milk'