Remove discovered mqtt_json light entity when discovery is cleared (#16906)

* Remove discovered mqtt_json entity device when discovery topic is cleared

* Keep imports ordered

homeassistant/components/light/mqtt_json.py

@@ -4,29 +4,31 @@ Support for MQTT JSON lights.
 For more details about this platform, please refer to the documentation at
 https://home-assistant.io/components/light.mqtt_json/
 """
-import logging
 import json
+import logging
+from typing import Optional
+
 import voluptuous as vol
 
-from homeassistant.core import callback
 from homeassistant.components import mqtt
 from homeassistant.components.light import (
-    ATTR_BRIGHTNESS, ATTR_COLOR_TEMP, ATTR_EFFECT, ATTR_FLASH,
-    ATTR_TRANSITION, ATTR_WHITE_VALUE, ATTR_HS_COLOR,
-    FLASH_LONG, FLASH_SHORT, Light, PLATFORM_SCHEMA, SUPPORT_BRIGHTNESS,
-    SUPPORT_COLOR_TEMP, SUPPORT_EFFECT, SUPPORT_FLASH, SUPPORT_COLOR,
-    SUPPORT_TRANSITION, SUPPORT_WHITE_VALUE)
+    ATTR_BRIGHTNESS, ATTR_COLOR_TEMP, ATTR_EFFECT, ATTR_FLASH, ATTR_HS_COLOR,
+    ATTR_TRANSITION, ATTR_WHITE_VALUE, FLASH_LONG, FLASH_SHORT,
+    PLATFORM_SCHEMA, SUPPORT_BRIGHTNESS, SUPPORT_COLOR, SUPPORT_COLOR_TEMP,
+    SUPPORT_EFFECT, SUPPORT_FLASH, SUPPORT_TRANSITION, SUPPORT_WHITE_VALUE,
+    Light)
 from homeassistant.components.light.mqtt import CONF_BRIGHTNESS_SCALE
-from homeassistant.const import (
-    CONF_BRIGHTNESS, CONF_COLOR_TEMP, CONF_EFFECT, STATE_ON,
-    CONF_NAME, CONF_OPTIMISTIC, CONF_RGB, CONF_WHITE_VALUE, CONF_XY)
 from homeassistant.components.mqtt import (
-    CONF_AVAILABILITY_TOPIC, CONF_STATE_TOPIC, CONF_COMMAND_TOPIC,
+    ATTR_DISCOVERY_HASH, CONF_AVAILABILITY_TOPIC, CONF_COMMAND_TOPIC,
     CONF_PAYLOAD_AVAILABLE, CONF_PAYLOAD_NOT_AVAILABLE, CONF_QOS, CONF_RETAIN,
-    MqttAvailability)
+    CONF_STATE_TOPIC, MqttAvailability, MqttDiscoveryUpdate)
+from homeassistant.const import (
+    CONF_BRIGHTNESS, CONF_COLOR_TEMP, CONF_EFFECT, CONF_NAME, CONF_OPTIMISTIC,
+    CONF_RGB, CONF_WHITE_VALUE, CONF_XY, STATE_ON)
+from homeassistant.core import callback
 import homeassistant.helpers.config_validation as cv
-from homeassistant.helpers.typing import HomeAssistantType, ConfigType
 from homeassistant.helpers.restore_state import async_get_last_state
+from homeassistant.helpers.typing import ConfigType, HomeAssistantType
 import homeassistant.util.color as color_util
 
 _LOGGER = logging.getLogger(__name__)
@@ -87,6 +89,11 @@ async def async_setup_platform(hass: HomeAssistantType, config: ConfigType,
     """Set up a MQTT JSON Light."""
     if discovery_info is not None:
         config = PLATFORM_SCHEMA(discovery_info)
+
+    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([MqttJson(
         config.get(CONF_NAME),
         config.get(CONF_UNIQUE_ID),
@@ -116,20 +123,23 @@ async def async_setup_platform(hass: HomeAssistantType, config: ConfigType,
         config.get(CONF_AVAILABILITY_TOPIC),
         config.get(CONF_PAYLOAD_AVAILABLE),
         config.get(CONF_PAYLOAD_NOT_AVAILABLE),
-        config.get(CONF_BRIGHTNESS_SCALE)
+        config.get(CONF_BRIGHTNESS_SCALE),
+        discovery_hash,
     )])
 
 
-class MqttJson(MqttAvailability, Light):
+class MqttJson(MqttAvailability, MqttDiscoveryUpdate, Light):
     """Representation of a MQTT JSON light."""
 
     def __init__(self, name, unique_id, effect_list, topic, qos, retain,
                  optimistic, brightness, color_temp, effect, rgb, white_value,
                  xy, hs, flash_times, availability_topic, payload_available,
-                 payload_not_available, brightness_scale):
+                 payload_not_available, brightness_scale,
+                 discovery_hash: Optional[str]):
         """Initialize MQTT JSON light."""
-        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._name = name
         self._unique_id = unique_id
         self._effect_list = effect_list
@@ -180,7 +190,8 @@ class MqttJson(MqttAvailability, Light):
 
     async def async_added_to_hass(self):
         """Subscribe to MQTT events."""
-        await super().async_added_to_hass()
+        await MqttAvailability.async_added_to_hass(self)
+        await MqttDiscoveryUpdate.async_added_to_hass(self)
 
         last_state = await async_get_last_state(self.hass, self.entity_id)
 

tests/components/light/test_mqtt_json.py

@@ -97,11 +97,12 @@ from homeassistant.const import (
     STATE_ON, STATE_OFF, STATE_UNAVAILABLE, ATTR_ASSUMED_STATE,
     ATTR_SUPPORTED_FEATURES)
 import homeassistant.components.light as light
+from homeassistant.components.mqtt.discovery import async_start
 import homeassistant.core as ha
 
 from tests.common import (
     get_test_home_assistant, mock_mqtt_component, fire_mqtt_message,
-    assert_setup_component, mock_coro)
+    assert_setup_component, mock_coro, async_fire_mqtt_message)
 from tests.components.light import common
 
 
@@ -669,3 +670,25 @@ class TestLightMQTTJSON(unittest.TestCase):
 
         state = self.hass.states.get('light.test')
         self.assertEqual(STATE_UNAVAILABLE, state.state)
+
+
+async def test_discovery_removal(hass, mqtt_mock, caplog):
+    """Test removal of discovered mqtt_json lights."""
+    await async_start(hass, 'homeassistant', {})
+    data = (
+        '{ "name": "Beer",'
+        '  "platform": "mqtt_json",'
+        '  "command_topic": "test_topic" }'
+    )
+    async_fire_mqtt_message(hass, 'homeassistant/light/bla/config',
+                            data)
+    await hass.async_block_till_done()
+    state = hass.states.get('light.beer')
+    assert state is not None
+    assert state.name == 'Beer'
+    async_fire_mqtt_message(hass, 'homeassistant/light/bla/config',
+                            '')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('light.beer')
+    assert state is None