[2/3] vacuum mqtt-discovery (#19478)

* add discoverability to mqtt-vacuum

homeassistant/components/mqtt/discovery.py

@@ -25,7 +25,7 @@ TOPIC_MATCHER = re.compile(
 SUPPORTED_COMPONENTS = [
     'binary_sensor', 'camera', 'cover', 'fan',
     'light', 'sensor', 'switch', 'lock', 'climate',
-    'alarm_control_panel']
+    'alarm_control_panel', 'vacuum']
 
 CONFIG_ENTRY_COMPONENTS = [
     'binary_sensor',
@@ -38,6 +38,7 @@ CONFIG_ENTRY_COMPONENTS = [
     'climate',
     'alarm_control_panel',
     'fan',
+    'vacuum',
 ]
 
 DEPRECATED_PLATFORM_TO_SCHEMA = {
@@ -69,12 +70,25 @@ ABBREVIATIONS = {
     'bri_stat_t': 'brightness_state_topic',
     'bri_val_tpl': 'brightness_value_template',
     'clr_temp_cmd_tpl': 'color_temp_command_template',
+    'bat_lev_t': 'battery_level_topic',
+    'bat_lev_tpl': 'battery_level_template',
+    'chrg_t': 'charging_topic',
+    'chrg_tpl': 'charging_template',
     'clr_temp_cmd_t': 'color_temp_command_topic',
     'clr_temp_stat_t': 'color_temp_state_topic',
     'clr_temp_val_tpl': 'color_temp_value_template',
+    'cln_t': 'cleaning_topic',
+    'cln_tpl': 'cleaning_template',
     'cmd_t': 'command_topic',
     'curr_temp_t': 'current_temperature_topic',
     'dev_cla': 'device_class',
+    'dock_t': 'docked_topic',
+    'dock_tpl': 'docked_template',
+    'err_t': 'error_topic',
+    'err_tpl': 'error_template',
+    'fanspd_t': 'fan_speed_topic',
+    'fanspd_tpl': 'fan_speed_template',
+    'fanspd_lst': 'fan_speed_list',
     'fx_cmd_t': 'effect_command_topic',
     'fx_list': 'effect_list',
     'fx_stat_t': 'effect_state_topic',
@@ -124,6 +138,7 @@ ABBREVIATIONS = {
     'rgb_cmd_t': 'rgb_command_topic',
     'rgb_stat_t': 'rgb_state_topic',
     'rgb_val_tpl': 'rgb_value_template',
+    'send_cmd_t': 'send_command_topic',
     'send_if_off': 'send_if_off',
     'set_pos_tpl': 'set_position_template',
     'set_pos_t': 'set_position_topic',
@@ -137,6 +152,7 @@ ABBREVIATIONS = {
     'stat_open': 'state_open',
     'stat_t': 'state_topic',
     'stat_val_tpl': 'state_value_template',
+    'sup_feat': 'supported_features',
     'swing_mode_cmd_t': 'swing_mode_command_topic',
     'swing_mode_stat_tpl': 'swing_mode_state_template',
     'swing_mode_stat_t': 'swing_mode_state_topic',

homeassistant/components/vacuum/__init__.py

@@ -95,7 +95,7 @@ def is_on(hass, entity_id=None):
 
 async def async_setup(hass, config):
     """Set up the vacuum component."""
-    component = EntityComponent(
+    component = hass.data[DOMAIN] = EntityComponent(
         _LOGGER, DOMAIN, hass, SCAN_INTERVAL, GROUP_NAME_ALL_VACUUMS)
 
     await component.async_setup(config)
@@ -152,6 +152,16 @@ async def async_setup(hass, config):
     return True
 
 
+async def async_setup_entry(hass, entry):
+    """Set up a config entry."""
+    return await hass.data[DOMAIN].async_setup_entry(entry)
+
+
+async def async_unload_entry(hass, entry):
+    """Unload a config entry."""
+    return await hass.data[DOMAIN].async_unload_entry(entry)
+
+
 class _BaseVacuum(Entity):
     """Representation of a base vacuum.
 

homeassistant/components/vacuum/mqtt.py

@@ -10,15 +10,18 @@ import voluptuous as vol
 
 from homeassistant.components import mqtt
 from homeassistant.components.mqtt import (
-    MqttAvailability, subscription)
+    ATTR_DISCOVERY_HASH, MqttAvailability, MqttDiscoveryUpdate,
+    subscription)
+from homeassistant.components.mqtt.discovery import MQTT_DISCOVERY_NEW
 from homeassistant.components.vacuum import (
     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)
+    VacuumDevice, DOMAIN)
 from homeassistant.const import ATTR_SUPPORTED_FEATURES, 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
 
 _LOGGER = logging.getLogger(__name__)
@@ -145,11 +148,30 @@ PLATFORM_SCHEMA = mqtt.MQTT_BASE_PLATFORM_SCHEMA.extend({
 
 async def async_setup_platform(hass, config, async_add_entities,
                                discovery_info=None):
-    """Set up the vacuum."""
+    """Set up MQTT vacuum through configuration.yaml."""
-    async_add_entities([MqttVacuum(config, discovery_info)])
+    await _async_setup_entity(config, async_add_entities,
+                              discovery_info)
 
 
-class MqttVacuum(MqttAvailability, VacuumDevice):
+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."""
+        config = PLATFORM_SCHEMA(discovery_payload)
+        await _async_setup_entity(config, async_add_entities,
+                                  discovery_payload[ATTR_DISCOVERY_HASH])
+
+    async_dispatcher_connect(
+        hass, MQTT_DISCOVERY_NEW.format(DOMAIN, 'mqtt'), async_discover)
+
+
+async def _async_setup_entity(config, async_add_entities,
+                              discovery_hash=None):
+    """Set up the MQTT vacuum."""
+    async_add_entities([MqttVacuum(config, discovery_hash)])
+
+
+class MqttVacuum(MqttAvailability, MqttDiscoveryUpdate, VacuumDevice):
     """Representation of a MQTT-controlled vacuum."""
 
     def __init__(self, config, discovery_info):
@@ -174,6 +196,8 @@ class MqttVacuum(MqttAvailability, VacuumDevice):
 
         MqttAvailability.__init__(self, availability_topic, qos,
                                   payload_available, payload_not_available)
+        MqttDiscoveryUpdate.__init__(self, discovery_info,
+                                     self.discovery_update)
 
     def _setup_from_config(self, config):
         self._name = config.get(CONF_NAME)
@@ -221,11 +245,24 @@ class MqttVacuum(MqttAvailability, VacuumDevice):
             )
         }
 
+    async def discovery_update(self, discovery_payload):
+        """Handle updated discovery message."""
+        config = PLATFORM_SCHEMA(discovery_payload)
+        self._setup_from_config(config)
+        await self.availability_discovery_update(config)
+        await self._subscribe_topics()
+        self.async_schedule_update_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 MqttAvailability.async_will_remove_from_hass(self)
+
     async def _subscribe_topics(self):
         """(Re)Subscribe to topics."""
         for tpl in self._templates.values():

tests/components/vacuum/common.py

@@ -10,92 +10,189 @@ from homeassistant.components.vacuum import (
 from homeassistant.const import (
     ATTR_COMMAND, ATTR_ENTITY_ID, SERVICE_TOGGLE,
     SERVICE_TURN_OFF, SERVICE_TURN_ON)
+from homeassistant.core import callback
 from homeassistant.loader import bind_hass
 
 
 @bind_hass
 def turn_on(hass, entity_id=None):
+    """Turn all or specified vacuum on."""
+    hass.add_job(async_turn_on, hass, entity_id)
+
+
+@callback
+@bind_hass
+def async_turn_on(hass, entity_id=None):
     """Turn all or specified vacuum on."""
     data = {ATTR_ENTITY_ID: entity_id} if entity_id else None
-    hass.services.call(DOMAIN, SERVICE_TURN_ON, data)
+    hass.async_add_job(hass.services.async_call(
+        DOMAIN, SERVICE_TURN_ON, data))
 
 
 @bind_hass
 def turn_off(hass, entity_id=None):
+    """Turn all or specified vacuum off."""
+    hass.add_job(async_turn_off, hass, entity_id)
+
+
+@callback
+@bind_hass
+def async_turn_off(hass, entity_id=None):
     """Turn all or specified vacuum off."""
     data = {ATTR_ENTITY_ID: entity_id} if entity_id else None
-    hass.services.call(DOMAIN, SERVICE_TURN_OFF, data)
+    hass.async_add_job(hass.services.async_call(
+        DOMAIN, SERVICE_TURN_OFF, data))
 
 
 @bind_hass
 def toggle(hass, entity_id=None):
+    """Toggle all or specified vacuum."""
+    hass.add_job(async_toggle, hass, entity_id)
+
+
+@callback
+@bind_hass
+def async_toggle(hass, entity_id=None):
     """Toggle all or specified vacuum."""
     data = {ATTR_ENTITY_ID: entity_id} if entity_id else None
-    hass.services.call(DOMAIN, SERVICE_TOGGLE, data)
+    hass.async_add_job(hass.services.async_call(
+        DOMAIN, SERVICE_TOGGLE, data))
 
 
 @bind_hass
 def locate(hass, entity_id=None):
+    """Locate all or specified vacuum."""
+    hass.add_job(async_locate, hass, entity_id)
+
+
+@callback
+@bind_hass
+def async_locate(hass, entity_id=None):
     """Locate all or specified vacuum."""
     data = {ATTR_ENTITY_ID: entity_id} if entity_id else None
-    hass.services.call(DOMAIN, SERVICE_LOCATE, data)
+    hass.async_add_job(hass.services.async_call(
+        DOMAIN, SERVICE_LOCATE, data))
 
 
 @bind_hass
 def clean_spot(hass, entity_id=None):
+    """Tell all or specified vacuum to perform a spot clean-up."""
+    hass.add_job(async_clean_spot, hass, entity_id)
+
+
+@callback
+@bind_hass
+def async_clean_spot(hass, entity_id=None):
     """Tell all or specified vacuum to perform a spot clean-up."""
     data = {ATTR_ENTITY_ID: entity_id} if entity_id else None
-    hass.services.call(DOMAIN, SERVICE_CLEAN_SPOT, data)
+    hass.async_add_job(hass.services.async_call(
+        DOMAIN, SERVICE_CLEAN_SPOT, data))
 
 
 @bind_hass
 def return_to_base(hass, entity_id=None):
+    """Tell all or specified vacuum to return to base."""
+    hass.add_job(async_return_to_base, hass, entity_id)
+
+
+@callback
+@bind_hass
+def async_return_to_base(hass, entity_id=None):
     """Tell all or specified vacuum to return to base."""
     data = {ATTR_ENTITY_ID: entity_id} if entity_id else None
-    hass.services.call(DOMAIN, SERVICE_RETURN_TO_BASE, data)
+    hass.async_add_job(hass.services.async_call(
+        DOMAIN, SERVICE_RETURN_TO_BASE, data))
 
 
 @bind_hass
 def start_pause(hass, entity_id=None):
+    """Tell all or specified vacuum to start or pause the current task."""
+    hass.add_job(async_start_pause, hass, entity_id)
+
+
+@callback
+@bind_hass
+def async_start_pause(hass, entity_id=None):
     """Tell all or specified vacuum to start or pause the current task."""
     data = {ATTR_ENTITY_ID: entity_id} if entity_id else None
-    hass.services.call(DOMAIN, SERVICE_START_PAUSE, data)
+    hass.async_add_job(hass.services.async_call(
+        DOMAIN, SERVICE_START_PAUSE, data))
 
 
 @bind_hass
 def start(hass, entity_id=None):
+    """Tell all or specified vacuum to start or resume the current task."""
+    hass.add_job(async_start, hass, entity_id)
+
+
+@callback
+@bind_hass
+def async_start(hass, entity_id=None):
     """Tell all or specified vacuum to start or resume the current task."""
     data = {ATTR_ENTITY_ID: entity_id} if entity_id else None
-    hass.services.call(DOMAIN, SERVICE_START, data)
+    hass.async_add_job(hass.services.async_call(
+        DOMAIN, SERVICE_START, data))
 
 
 @bind_hass
 def pause(hass, entity_id=None):
+    """Tell all or the specified vacuum to pause the current task."""
+    hass.add_job(async_pause, hass, entity_id)
+
+
+@callback
+@bind_hass
+def async_pause(hass, entity_id=None):
     """Tell all or the specified vacuum to pause the current task."""
     data = {ATTR_ENTITY_ID: entity_id} if entity_id else None
-    hass.services.call(DOMAIN, SERVICE_PAUSE, data)
+    hass.async_add_job(hass.services.async_call(
+        DOMAIN, SERVICE_PAUSE, data))
 
 
 @bind_hass
 def stop(hass, entity_id=None):
+    """Stop all or specified vacuum."""
+    hass.add_job(async_stop, hass, entity_id)
+
+
+@callback
+@bind_hass
+def async_stop(hass, entity_id=None):
     """Stop all or specified vacuum."""
     data = {ATTR_ENTITY_ID: entity_id} if entity_id else None
-    hass.services.call(DOMAIN, SERVICE_STOP, data)
+    hass.async_add_job(hass.services.async_call(
+        DOMAIN, SERVICE_STOP, data))
 
 
 @bind_hass
 def set_fan_speed(hass, fan_speed, entity_id=None):
+    """Set fan speed for all or specified vacuum."""
+    hass.add_job(async_set_fan_speed, hass, fan_speed, entity_id)
+
+
+@callback
+@bind_hass
+def async_set_fan_speed(hass, fan_speed, entity_id=None):
     """Set fan speed for all or specified vacuum."""
     data = {ATTR_ENTITY_ID: entity_id} if entity_id else {}
     data[ATTR_FAN_SPEED] = fan_speed
-    hass.services.call(DOMAIN, SERVICE_SET_FAN_SPEED, data)
+    hass.async_add_job(hass.services.async_call(
+        DOMAIN, SERVICE_SET_FAN_SPEED, data))
 
 
 @bind_hass
 def send_command(hass, command, params=None, entity_id=None):
+    """Send command to all or specified vacuum."""
+    hass.add_job(async_send_command, hass, command, params, entity_id)
+
+
+@callback
+@bind_hass
+def async_send_command(hass, command, params=None, entity_id=None):
     """Send command to all or specified vacuum."""
     data = {ATTR_ENTITY_ID: entity_id} if entity_id else {}
     data[ATTR_COMMAND] = command
     if params is not None:
         data[ATTR_PARAMS] = params
-    hass.services.call(DOMAIN, SERVICE_SEND_COMMAND, data)
+    hass.async_add_job(hass.services.async_call(
+        DOMAIN, SERVICE_SEND_COMMAND, data))

tests/components/vacuum/test_mqtt.py

@@ -1,136 +1,143 @@
-"""The tests for the Demo vacuum platform."""
+"""The tests for the Mqtt vacuum platform."""
-import unittest
+import pytest
 
-from homeassistant.components import vacuum
+from homeassistant.setup import async_setup_component
-from homeassistant.components.vacuum import (
-    ATTR_BATTERY_LEVEL, ATTR_BATTERY_ICON, ATTR_STATUS,
-    ATTR_FAN_SPEED, mqtt)
-from homeassistant.components.mqtt import CONF_COMMAND_TOPIC
 from homeassistant.const import (
     CONF_PLATFORM, STATE_OFF, STATE_ON, STATE_UNAVAILABLE, CONF_NAME)
-from homeassistant.setup import setup_component
+from homeassistant.components import vacuum, mqtt
-
+from homeassistant.components.vacuum import (
+    ATTR_BATTERY_LEVEL, ATTR_BATTERY_ICON, ATTR_STATUS,
+    ATTR_FAN_SPEED, mqtt as mqttvacuum)
+from homeassistant.components.mqtt import CONF_COMMAND_TOPIC
+from homeassistant.components.mqtt.discovery import async_start
 from tests.common import (
-    fire_mqtt_message, get_test_home_assistant, mock_mqtt_component)
+    async_mock_mqtt_component,
+    async_fire_mqtt_message, MockConfigEntry)
 from tests.components.vacuum import common
 
 
-class TestVacuumMQTT(unittest.TestCase):
+default_config = {
-    """MQTT vacuum component test class."""
-
-    def setUp(self):  # pylint: disable=invalid-name
-        """Set up things to be run when tests are started."""
-        self.hass = get_test_home_assistant()
-        self.mock_publish = mock_mqtt_component(self.hass)
-
-        self.default_config = {
     CONF_PLATFORM: 'mqtt',
     CONF_NAME: 'mqtttest',
     CONF_COMMAND_TOPIC: 'vacuum/command',
-            mqtt.CONF_SEND_COMMAND_TOPIC: 'vacuum/send_command',
+    mqttvacuum.CONF_SEND_COMMAND_TOPIC: 'vacuum/send_command',
-            mqtt.CONF_BATTERY_LEVEL_TOPIC: 'vacuum/state',
+    mqttvacuum.CONF_BATTERY_LEVEL_TOPIC: 'vacuum/state',
-            mqtt.CONF_BATTERY_LEVEL_TEMPLATE:
+    mqttvacuum.CONF_BATTERY_LEVEL_TEMPLATE:
         '{{ value_json.battery_level }}',
-            mqtt.CONF_CHARGING_TOPIC: 'vacuum/state',
+    mqttvacuum.CONF_CHARGING_TOPIC: 'vacuum/state',
-            mqtt.CONF_CHARGING_TEMPLATE: '{{ value_json.charging }}',
+    mqttvacuum.CONF_CHARGING_TEMPLATE: '{{ value_json.charging }}',
-            mqtt.CONF_CLEANING_TOPIC: 'vacuum/state',
+    mqttvacuum.CONF_CLEANING_TOPIC: 'vacuum/state',
-            mqtt.CONF_CLEANING_TEMPLATE: '{{ value_json.cleaning }}',
+    mqttvacuum.CONF_CLEANING_TEMPLATE: '{{ value_json.cleaning }}',
-            mqtt.CONF_DOCKED_TOPIC: 'vacuum/state',
+    mqttvacuum.CONF_DOCKED_TOPIC: 'vacuum/state',
-            mqtt.CONF_DOCKED_TEMPLATE: '{{ value_json.docked }}',
+    mqttvacuum.CONF_DOCKED_TEMPLATE: '{{ value_json.docked }}',
-            mqtt.CONF_STATE_TOPIC: 'vacuum/state',
+    mqttvacuum.CONF_STATE_TOPIC: 'vacuum/state',
-            mqtt.CONF_STATE_TEMPLATE: '{{ value_json.state }}',
+    mqttvacuum.CONF_STATE_TEMPLATE: '{{ value_json.state }}',
-            mqtt.CONF_FAN_SPEED_TOPIC: 'vacuum/state',
+    mqttvacuum.CONF_FAN_SPEED_TOPIC: 'vacuum/state',
-            mqtt.CONF_FAN_SPEED_TEMPLATE: '{{ value_json.fan_speed }}',
+    mqttvacuum.CONF_FAN_SPEED_TEMPLATE: '{{ value_json.fan_speed }}',
-            mqtt.CONF_SET_FAN_SPEED_TOPIC: 'vacuum/set_fan_speed',
+    mqttvacuum.CONF_SET_FAN_SPEED_TOPIC: 'vacuum/set_fan_speed',
-            mqtt.CONF_FAN_SPEED_LIST: ['min', 'medium', 'high', 'max'],
+    mqttvacuum.CONF_FAN_SPEED_LIST: ['min', 'medium', 'high', 'max'],
 }
 
-    def tearDown(self):  # pylint: disable=invalid-name
-        """Stop down everything that was started."""
-        self.hass.stop()
 
-    def test_default_supported_features(self):
+@pytest.fixture
+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 setup_component(self.hass, vacuum.DOMAIN, {
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
-            vacuum.DOMAIN: self.default_config,
+        vacuum.DOMAIN: default_config,
     })
-        entity = self.hass.states.get('vacuum.mqtttest')
+    entity = hass.states.get('vacuum.mqtttest')
     entity_features = \
-            entity.attributes.get(mqtt.CONF_SUPPORTED_FEATURES, 0)
+        entity.attributes.get(mqttvacuum.CONF_SUPPORTED_FEATURES, 0)
-        assert sorted(mqtt.services_to_strings(entity_features)) == \
+    assert sorted(mqttvacuum.services_to_strings(entity_features)) == \
         sorted(['turn_on', 'turn_off', 'stop',
                 'return_home', 'battery', 'status',
                 'clean_spot'])
 
-    def test_all_commands(self):
-        """Test simple commands to the vacuum."""
-        self.default_config[mqtt.CONF_SUPPORTED_FEATURES] = \
-            mqtt.services_to_strings(mqtt.ALL_SERVICES)
 
-        assert setup_component(self.hass, vacuum.DOMAIN, {
+async def test_all_commands(hass, mock_publish):
-            vacuum.DOMAIN: self.default_config,
+    """Test simple commands to the vacuum."""
+    default_config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES)
+
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: default_config,
     })
 
-        common.turn_on(self.hass, 'vacuum.mqtttest')
+    common.turn_on(hass, 'vacuum.mqtttest')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.mock_publish.async_publish.assert_called_once_with(
+    await hass.async_block_till_done()
+    mock_publish.async_publish.assert_called_once_with(
         'vacuum/command', 'turn_on', 0, False)
-        self.mock_publish.async_publish.reset_mock()
+    mock_publish.async_publish.reset_mock()
 
-        common.turn_off(self.hass, 'vacuum.mqtttest')
+    common.turn_off(hass, 'vacuum.mqtttest')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.mock_publish.async_publish.assert_called_once_with(
+    await hass.async_block_till_done()
+    mock_publish.async_publish.assert_called_once_with(
         'vacuum/command', 'turn_off', 0, False)
-        self.mock_publish.async_publish.reset_mock()
+    mock_publish.async_publish.reset_mock()
 
-        common.stop(self.hass, 'vacuum.mqtttest')
+    common.stop(hass, 'vacuum.mqtttest')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.mock_publish.async_publish.assert_called_once_with(
+    await hass.async_block_till_done()
+    mock_publish.async_publish.assert_called_once_with(
         'vacuum/command', 'stop', 0, False)
-        self.mock_publish.async_publish.reset_mock()
+    mock_publish.async_publish.reset_mock()
 
-        common.clean_spot(self.hass, 'vacuum.mqtttest')
+    common.clean_spot(hass, 'vacuum.mqtttest')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.mock_publish.async_publish.assert_called_once_with(
+    await hass.async_block_till_done()
+    mock_publish.async_publish.assert_called_once_with(
         'vacuum/command', 'clean_spot', 0, False)
-        self.mock_publish.async_publish.reset_mock()
+    mock_publish.async_publish.reset_mock()
 
-        common.locate(self.hass, 'vacuum.mqtttest')
+    common.locate(hass, 'vacuum.mqtttest')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.mock_publish.async_publish.assert_called_once_with(
+    await hass.async_block_till_done()
+    mock_publish.async_publish.assert_called_once_with(
         'vacuum/command', 'locate', 0, False)
-        self.mock_publish.async_publish.reset_mock()
+    mock_publish.async_publish.reset_mock()
 
-        common.start_pause(self.hass, 'vacuum.mqtttest')
+    common.start_pause(hass, 'vacuum.mqtttest')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.mock_publish.async_publish.assert_called_once_with(
+    await hass.async_block_till_done()
+    mock_publish.async_publish.assert_called_once_with(
         'vacuum/command', 'start_pause', 0, False)
-        self.mock_publish.async_publish.reset_mock()
+    mock_publish.async_publish.reset_mock()
 
-        common.return_to_base(self.hass, 'vacuum.mqtttest')
+    common.return_to_base(hass, 'vacuum.mqtttest')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.mock_publish.async_publish.assert_called_once_with(
+    await hass.async_block_till_done()
+    mock_publish.async_publish.assert_called_once_with(
         'vacuum/command', 'return_to_base', 0, False)
-        self.mock_publish.async_publish.reset_mock()
+    mock_publish.async_publish.reset_mock()
 
-        common.set_fan_speed(self.hass, 'high', 'vacuum.mqtttest')
+    common.set_fan_speed(hass, 'high', 'vacuum.mqtttest')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.mock_publish.async_publish.assert_called_once_with(
+    await hass.async_block_till_done()
+    mock_publish.async_publish.assert_called_once_with(
         'vacuum/set_fan_speed', 'high', 0, False)
-        self.mock_publish.async_publish.reset_mock()
+    mock_publish.async_publish.reset_mock()
 
-        common.send_command(self.hass, '44 FE 93', entity_id='vacuum.mqtttest')
+    common.send_command(hass, '44 FE 93', entity_id='vacuum.mqtttest')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.mock_publish.async_publish.assert_called_once_with(
+    await hass.async_block_till_done()
+    mock_publish.async_publish.assert_called_once_with(
         'vacuum/send_command', '44 FE 93', 0, False)
 
-    def test_status(self):
-        """Test status updates from the vacuum."""
-        self.default_config[mqtt.CONF_SUPPORTED_FEATURES] = \
-            mqtt.services_to_strings(mqtt.ALL_SERVICES)
 
-        assert setup_component(self.hass, vacuum.DOMAIN, {
+async def test_status(hass, mock_publish):
-            vacuum.DOMAIN: self.default_config,
+    """Test status updates from the vacuum."""
+    default_config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES)
+
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: default_config,
     })
 
     message = """{
@@ -140,9 +147,10 @@ class TestVacuumMQTT(unittest.TestCase):
         "charging": false,
         "fan_speed": "max"
     }"""
-        fire_mqtt_message(self.hass, 'vacuum/state', message)
+    async_fire_mqtt_message(hass, 'vacuum/state', message)
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        state = self.hass.states.get('vacuum.mqtttest')
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.mqtttest')
     assert STATE_ON == state.state
     assert 'mdi:battery-50' == \
         state.attributes.get(ATTR_BATTERY_ICON)
@@ -157,98 +165,168 @@ class TestVacuumMQTT(unittest.TestCase):
         "fan_speed": "min"
     }"""
 
-        fire_mqtt_message(self.hass, 'vacuum/state', message)
+    async_fire_mqtt_message(hass, 'vacuum/state', message)
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        state = self.hass.states.get('vacuum.mqtttest')
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.mqtttest')
     assert STATE_OFF == state.state
     assert 'mdi:battery-charging-60' == \
         state.attributes.get(ATTR_BATTERY_ICON)
     assert 61 == state.attributes.get(ATTR_BATTERY_LEVEL)
     assert 'min' == state.attributes.get(ATTR_FAN_SPEED)
 
-    def test_battery_template(self):
+
+async def test_battery_template(hass, mock_publish):
     """Test that you can use non-default templates for battery_level."""
-        self.default_config.update({
+    default_config.update({
-            mqtt.CONF_SUPPORTED_FEATURES:
+        mqttvacuum.CONF_SUPPORTED_FEATURES:
-                mqtt.services_to_strings(mqtt.ALL_SERVICES),
+            mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES),
-            mqtt.CONF_BATTERY_LEVEL_TOPIC: "retroroomba/battery_level",
+        mqttvacuum.CONF_BATTERY_LEVEL_TOPIC: "retroroomba/battery_level",
-            mqtt.CONF_BATTERY_LEVEL_TEMPLATE: "{{ value }}"
+        mqttvacuum.CONF_BATTERY_LEVEL_TEMPLATE: "{{ value }}"
     })
 
-        assert setup_component(self.hass, vacuum.DOMAIN, {
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
-            vacuum.DOMAIN: self.default_config,
+        vacuum.DOMAIN: default_config,
     })
 
-        fire_mqtt_message(self.hass, 'retroroomba/battery_level', '54')
+    async_fire_mqtt_message(hass, 'retroroomba/battery_level', '54')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        state = self.hass.states.get('vacuum.mqtttest')
+    state = hass.states.get('vacuum.mqtttest')
     assert 54 == state.attributes.get(ATTR_BATTERY_LEVEL)
     assert state.attributes.get(ATTR_BATTERY_ICON) == \
         'mdi:battery-50'
 
-    def test_status_invalid_json(self):
-        """Test to make sure nothing breaks if the vacuum sends bad JSON."""
-        self.default_config[mqtt.CONF_SUPPORTED_FEATURES] = \
-            mqtt.services_to_strings(mqtt.ALL_SERVICES)
 
-        assert setup_component(self.hass, vacuum.DOMAIN, {
+async def test_status_invalid_json(hass, mock_publish):
-            vacuum.DOMAIN: self.default_config,
+    """Test to make sure nothing breaks if the vacuum sends bad JSON."""
+    default_config[mqttvacuum.CONF_SUPPORTED_FEATURES] = \
+        mqttvacuum.services_to_strings(mqttvacuum.ALL_SERVICES)
+
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: default_config,
     })
 
-        fire_mqtt_message(self.hass, 'vacuum/state', '{"asdfasas false}')
+    async_fire_mqtt_message(hass, 'vacuum/state', '{"asdfasas false}')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        state = self.hass.states.get('vacuum.mqtttest')
+    state = hass.states.get('vacuum.mqtttest')
     assert STATE_OFF == state.state
     assert "Stopped" == state.attributes.get(ATTR_STATUS)
 
-    def test_default_availability_payload(self):
+
+async def test_default_availability_payload(hass, mock_publish):
     """Test availability by default payload with defined topic."""
-        self.default_config.update({
+    default_config.update({
         'availability_topic': 'availability-topic'
     })
 
-        assert setup_component(self.hass, vacuum.DOMAIN, {
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
-            vacuum.DOMAIN: self.default_config,
+        vacuum.DOMAIN: default_config,
     })
 
-        state = self.hass.states.get('vacuum.mqtttest')
+    state = hass.states.get('vacuum.mqtttest')
     assert STATE_UNAVAILABLE == state.state
 
-        fire_mqtt_message(self.hass, 'availability-topic', 'online')
+    async_fire_mqtt_message(hass, 'availability-topic', 'online')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get('vacuum.mqtttest')
+    state = hass.states.get('vacuum.mqtttest')
     assert STATE_UNAVAILABLE != state.state
 
-        fire_mqtt_message(self.hass, 'availability-topic', 'offline')
+    async_fire_mqtt_message(hass, 'availability-topic', 'offline')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get('vacuum.mqtttest')
+    state = hass.states.get('vacuum.mqtttest')
     assert STATE_UNAVAILABLE == state.state
 
-    def test_custom_availability_payload(self):
+
+async def test_custom_availability_payload(hass, mock_publish):
     """Test availability by custom payload with defined topic."""
-        self.default_config.update({
+    default_config.update({
         'availability_topic': 'availability-topic',
         'payload_available': 'good',
         'payload_not_available': 'nogood'
     })
 
-        assert setup_component(self.hass, vacuum.DOMAIN, {
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
-            vacuum.DOMAIN: self.default_config,
+        vacuum.DOMAIN: default_config,
     })
 
-        state = self.hass.states.get('vacuum.mqtttest')
+    state = hass.states.get('vacuum.mqtttest')
     assert STATE_UNAVAILABLE == state.state
 
-        fire_mqtt_message(self.hass, 'availability-topic', 'good')
+    async_fire_mqtt_message(hass, 'availability-topic', 'good')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get('vacuum.mqtttest')
+    state = hass.states.get('vacuum.mqtttest')
     assert STATE_UNAVAILABLE != state.state
 
-        fire_mqtt_message(self.hass, 'availability-topic', 'nogood')
+    async_fire_mqtt_message(hass, 'availability-topic', 'nogood')
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get('vacuum.mqtttest')
+    state = hass.states.get('vacuum.mqtttest')
     assert STATE_UNAVAILABLE == state.state
+
+
+async def test_discovery_removal_vacuum(hass, mock_publish):
+    """Test removal of discovered vacuum."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    await async_start(hass, 'homeassistant', {}, entry)
+
+    data = (
+        '{ "name": "Beer",'
+        '  "command_topic": "test_topic" }'
+    )
+
+    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_update_vacuum(hass, mock_publish):
+    """Test update of discovered vacuum."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    await async_start(hass, 'homeassistant', {}, entry)
+
+    data1 = (
+        '{ "name": "Beer",'
+        '  "command_topic": "test_topic" }'
+    )
+    data2 = (
+        '{ "name": "Milk",'
+        '  "command_topic": "test_topic" }'
+    )
+
+    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