Remove discovered MQTT lock device when discovery topic is cleared (#16859)

homeassistant/components/lock/mqtt.py

@@ -12,9 +12,9 @@ import voluptuous as vol
 from homeassistant.core import callback
 from homeassistant.components.lock import LockDevice
 from homeassistant.components.mqtt import (
-    CONF_AVAILABILITY_TOPIC, CONF_STATE_TOPIC, CONF_COMMAND_TOPIC,
-    CONF_PAYLOAD_AVAILABLE, CONF_PAYLOAD_NOT_AVAILABLE, CONF_QOS, CONF_RETAIN,
-    MqttAvailability)
+    ATTR_DISCOVERY_HASH, CONF_AVAILABILITY_TOPIC, CONF_STATE_TOPIC,
+    CONF_COMMAND_TOPIC, CONF_PAYLOAD_AVAILABLE, CONF_PAYLOAD_NOT_AVAILABLE,
+    CONF_QOS, CONF_RETAIN, MqttAvailability, MqttDiscoveryUpdate)
 from homeassistant.const import (
     CONF_NAME, CONF_OPTIMISTIC, CONF_VALUE_TEMPLATE)
 from homeassistant.components import mqtt
@@ -52,6 +52,10 @@ def async_setup_platform(hass, config, async_add_entities,
     if value_template is not None:
         value_template.hass = hass
 
+    discovery_hash = None
+    if discovery_info is not None and ATTR_DISCOVERY_HASH in discovery_info:
+        discovery_hash = discovery_info[ATTR_DISCOVERY_HASH]
+
     async_add_entities([MqttLock(
         config.get(CONF_NAME),
         config.get(CONF_STATE_TOPIC),
@@ -64,19 +68,22 @@ def async_setup_platform(hass, config, async_add_entities,
         value_template,
         config.get(CONF_AVAILABILITY_TOPIC),
         config.get(CONF_PAYLOAD_AVAILABLE),
-        config.get(CONF_PAYLOAD_NOT_AVAILABLE)
+        config.get(CONF_PAYLOAD_NOT_AVAILABLE),
+        discovery_hash,
     )])
 
 
-class MqttLock(MqttAvailability, LockDevice):
+class MqttLock(MqttAvailability, MqttDiscoveryUpdate, LockDevice):
     """Representation of a lock that can be toggled using MQTT."""
 
     def __init__(self, name, state_topic, command_topic, qos, retain,
                  payload_lock, payload_unlock, optimistic, value_template,
-                 availability_topic, payload_available, payload_not_available):
+                 availability_topic, payload_available, payload_not_available,
+                 discovery_hash):
         """Initialize the lock."""
-        super().__init__(availability_topic, qos, payload_available,
-                         payload_not_available)
+        MqttAvailability.__init__(self, availability_topic, qos,
+                                  payload_available, payload_not_available)
+        MqttDiscoveryUpdate.__init__(self, discovery_hash)
         self._state = False
         self._name = name
         self._state_topic = state_topic
@@ -87,11 +94,13 @@ class MqttLock(MqttAvailability, LockDevice):
         self._payload_unlock = payload_unlock
         self._optimistic = optimistic
         self._template = value_template
+        self._discovery_hash = discovery_hash
 
     @asyncio.coroutine
     def async_added_to_hass(self):
         """Subscribe to MQTT events."""
-        yield from super().async_added_to_hass()
+        yield from MqttAvailability.async_added_to_hass(self)
+        yield from MqttDiscoveryUpdate.async_added_to_hass(self)
 
         @callback
         def message_received(topic, payload, qos):

tests/components/lock/test_mqtt.py

@@ -5,8 +5,10 @@ from homeassistant.setup import setup_component
 from homeassistant.const import (
     STATE_LOCKED, STATE_UNLOCKED, STATE_UNAVAILABLE, ATTR_ASSUMED_STATE)
 import homeassistant.components.lock as lock
+from homeassistant.components.mqtt.discovery import async_start
 from tests.common import (
-    mock_mqtt_component, fire_mqtt_message, get_test_home_assistant)
+    mock_mqtt_component, async_fire_mqtt_message, fire_mqtt_message,
+    get_test_home_assistant)
 
 
 class TestLockMQTT(unittest.TestCase):
@@ -172,3 +174,24 @@ class TestLockMQTT(unittest.TestCase):
 
         state = self.hass.states.get('lock.test')
         self.assertEqual(STATE_UNAVAILABLE, state.state)
+
+
+async def test_discovery_removal_lock(hass, mqtt_mock, caplog):
+    """Test removal of discovered lock."""
+    await async_start(hass, 'homeassistant', {})
+    data = (
+        '{ "name": "Beer",'
+        '  "command_topic": "test_topic" }'
+    )
+    async_fire_mqtt_message(hass, 'homeassistant/lock/bla/config',
+                            data)
+    await hass.async_block_till_done()
+    state = hass.states.get('lock.beer')
+    assert state is not None
+    assert state.name == 'Beer'
+    async_fire_mqtt_message(hass, 'homeassistant/lock/bla/config',
+                            '')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('lock.beer')
+    assert state is None