Cleanup if discovered mqtt light can't be added (#19740)

* Cleanup if discovered mqtt light can't be added

* No bare except

* Clear ALREADY_DISCOVERED list with helper

* Use constant instead of string literal

homeassistant/components/mqtt/light/__init__.py

@@ -10,7 +10,8 @@ import voluptuous as vol
 
 from homeassistant.components import light
 from homeassistant.components.mqtt import ATTR_DISCOVERY_HASH
-from homeassistant.components.mqtt.discovery import MQTT_DISCOVERY_NEW
+from homeassistant.components.mqtt.discovery import (
+    MQTT_DISCOVERY_NEW, clear_discovery_hash)
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.typing import HomeAssistantType, ConfigType
 
@@ -44,23 +45,29 @@ PLATFORM_SCHEMA = vol.All(vol.Schema({
 async def async_setup_platform(hass: HomeAssistantType, config: ConfigType,
                                async_add_entities, discovery_info=None):
     """Set up MQTT light through configuration.yaml."""
-    await _async_setup_entity(hass, config, async_add_entities)
+    await _async_setup_entity(config, async_add_entities)
 
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up MQTT light dynamically through MQTT discovery."""
     async def async_discover(discovery_payload):
         """Discover and add a MQTT light."""
-        config = PLATFORM_SCHEMA(discovery_payload)
-        await _async_setup_entity(hass, config, async_add_entities,
-                                  discovery_payload[ATTR_DISCOVERY_HASH])
+        try:
+            discovery_hash = discovery_payload[ATTR_DISCOVERY_HASH]
+            config = PLATFORM_SCHEMA(discovery_payload)
+            await _async_setup_entity(config, async_add_entities,
+                                      discovery_hash)
+        except Exception:
+            if discovery_hash:
+                clear_discovery_hash(hass, discovery_hash)
+            raise
 
     async_dispatcher_connect(
         hass, MQTT_DISCOVERY_NEW.format(light.DOMAIN, 'mqtt'),
         async_discover)
 
 
-async def _async_setup_entity(hass, config, async_add_entities,
+async def _async_setup_entity(config, async_add_entities,
                               discovery_hash=None):
     """Set up a MQTT Light."""
     setup_entity = {
@@ -68,5 +75,5 @@ async def _async_setup_entity(hass, config, async_add_entities,
         'json': schema_json.async_setup_entity_json,
         'template': schema_template.async_setup_entity_template,
     }
-    await setup_entity[config['schema']](
-        hass, config, async_add_entities, discovery_hash)
+    await setup_entity[config[CONF_SCHEMA]](
+        config, async_add_entities, discovery_hash)

homeassistant/components/mqtt/light/schema_basic.py

@@ -112,7 +112,7 @@ PLATFORM_SCHEMA_BASIC = mqtt.MQTT_RW_PLATFORM_SCHEMA.extend({
     mqtt.MQTT_JSON_ATTRS_SCHEMA.schema)
 
 
-async def async_setup_entity_basic(hass, config, async_add_entities,
+async def async_setup_entity_basic(config, async_add_entities,
                                    discovery_hash=None):
     """Set up a MQTT Light."""
     config.setdefault(

homeassistant/components/mqtt/light/schema_json.py

@@ -25,7 +25,7 @@ from homeassistant.const import (
 from homeassistant.core import callback
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.restore_state import RestoreEntity
-from homeassistant.helpers.typing import ConfigType, HomeAssistantType
+from homeassistant.helpers.typing import ConfigType
 import homeassistant.util.color as color_util
 
 from .schema_basic import CONF_BRIGHTNESS_SCALE
@@ -85,7 +85,7 @@ PLATFORM_SCHEMA_JSON = mqtt.MQTT_RW_PLATFORM_SCHEMA.extend({
     mqtt.MQTT_JSON_ATTRS_SCHEMA.schema)
 
 
-async def async_setup_entity_json(hass: HomeAssistantType, config: ConfigType,
+async def async_setup_entity_json(config: ConfigType,
                                   async_add_entities, discovery_hash):
     """Set up a MQTT JSON Light."""
     async_add_entities([MqttLightJson(config, discovery_hash)])

homeassistant/components/mqtt/light/schema_template.py

@@ -71,7 +71,7 @@ PLATFORM_SCHEMA_TEMPLATE = mqtt.MQTT_RW_PLATFORM_SCHEMA.extend({
     mqtt.MQTT_JSON_ATTRS_SCHEMA.schema)
 
 
-async def async_setup_entity_template(hass, config, async_add_entities,
+async def async_setup_entity_template(config, async_add_entities,
                                       discovery_hash):
     """Set up a MQTT Template light."""
     async_add_entities([MqttTemplate(config, discovery_hash)])

tests/components/mqtt/test_light.py

@@ -1239,7 +1239,7 @@ async def test_discovery_deprecated(hass, mqtt_mock, caplog):
 
 
 async def test_discovery_update_light(hass, mqtt_mock, caplog):
-    """Test removal of discovered light."""
+    """Test update of discovered light."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)
     await async_start(hass, 'homeassistant', {}, entry)
 
@@ -1274,6 +1274,39 @@ async def test_discovery_update_light(hass, mqtt_mock, caplog):
     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" }'
+    )
+    data2 = (
+        '{ "name": "Milk",'
+        '  "status_topic": "test_topic",'
+        '  "command_topic": "test_topic" }'
+    )
+
+    async_fire_mqtt_message(hass, 'homeassistant/light/bla/config',
+                            data1)
+    await hass.async_block_till_done()
+
+    state = hass.states.get('light.beer')
+    assert state is None
+
+    async_fire_mqtt_message(hass, 'homeassistant/light/bla/config',
+                            data2)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('light.milk')
+    assert state is not None
+    assert state.name == 'Milk'
+    state = hass.states.get('light.beer')
+    assert state is None
+
+
 async def test_entity_device_info_with_identifier(hass, mqtt_mock):
     """Test MQTT light device registry integration."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)

tests/components/mqtt/test_light_json.py

@@ -707,7 +707,7 @@ async def test_discovery_deprecated(hass, mqtt_mock, caplog):
 
 
 async def test_discovery_update_light(hass, mqtt_mock, caplog):
-    """Test removal of discovered light."""
+    """Test update of discovered light."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)
     await async_start(hass, 'homeassistant', {}, entry)
 
@@ -744,6 +744,40 @@ async def test_discovery_update_light(hass, mqtt_mock, caplog):
     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" }'
+    )
+    data2 = (
+        '{ "name": "Milk",'
+        '  "schema": "json",'
+        '  "status_topic": "test_topic",'
+        '  "command_topic": "test_topic" }'
+    )
+
+    async_fire_mqtt_message(hass, 'homeassistant/light/bla/config',
+                            data1)
+    await hass.async_block_till_done()
+
+    state = hass.states.get('light.beer')
+    assert state is None
+
+    async_fire_mqtt_message(hass, 'homeassistant/light/bla/config',
+                            data2)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('light.milk')
+    assert state is not None
+    assert state.name == 'Milk'
+    state = hass.states.get('light.beer')
+    assert state is None
+
+
 async def test_entity_device_info_with_identifier(hass, mqtt_mock):
     """Test MQTT light device registry integration."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)

tests/components/mqtt/test_light_template.py

@@ -627,7 +627,7 @@ async def test_unique_id(hass):
     assert len(hass.states.async_entity_ids(light.DOMAIN)) == 1
 
 
-async def test_discovery(hass, mqtt_mock, caplog):
+async def test_discovery_removal(hass, mqtt_mock, caplog):
     """Test removal of discovered mqtt_json lights."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)
     await async_start(hass, 'homeassistant', {'mqtt': {}}, entry)
@@ -644,6 +644,12 @@ async def test_discovery(hass, mqtt_mock, caplog):
     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
 
 
 async def test_discovery_deprecated(hass, mqtt_mock, caplog):
@@ -666,7 +672,7 @@ async def test_discovery_deprecated(hass, mqtt_mock, caplog):
 
 
 async def test_discovery_update_light(hass, mqtt_mock, caplog):
-    """Test removal of discovered light."""
+    """Test update of discovered light."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)
     await async_start(hass, 'homeassistant', {}, entry)
 
@@ -707,6 +713,42 @@ async def test_discovery_update_light(hass, mqtt_mock, caplog):
     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" }'
+    )
+    data2 = (
+        '{ "name": "Milk",'
+        '  "schema": "template",'
+        '  "status_topic": "test_topic",'
+        '  "command_topic": "test_topic",'
+        '  "command_on_template": "on",'
+        '  "command_off_template": "off"}'
+    )
+
+    async_fire_mqtt_message(hass, 'homeassistant/light/bla/config',
+                            data1)
+    await hass.async_block_till_done()
+
+    state = hass.states.get('light.beer')
+    assert state is None
+
+    async_fire_mqtt_message(hass, 'homeassistant/light/bla/config',
+                            data2)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('light.milk')
+    assert state is not None
+    assert state.name == 'Milk'
+    state = hass.states.get('light.beer')
+    assert state is None
+
+
 async def test_entity_device_info_with_identifier(hass, mqtt_mock):
     """Test MQTT light device registry integration."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)