Cleanup if discovered mqtt switch can't be added (#19721)

* Cleanup if discovered mqtt switch can't be added

homeassistant/components/mqtt/__init__.py

@@ -972,7 +972,7 @@ class MqttDiscoveryUpdate(Entity):
 
         from homeassistant.helpers.dispatcher import async_dispatcher_connect
         from homeassistant.components.mqtt.discovery import (
-            ALREADY_DISCOVERED, MQTT_DISCOVERY_UPDATED)
+            MQTT_DISCOVERY_UPDATED, clear_discovery_hash)
 
         @callback
         def discovery_callback(payload):
@@ -983,7 +983,7 @@ class MqttDiscoveryUpdate(Entity):
                 # Empty payload: Remove component
                 _LOGGER.info("Removing component: %s", self.entity_id)
                 self.hass.async_create_task(self.async_remove())
-                del self.hass.data[ALREADY_DISCOVERED][self._discovery_hash]
+                clear_discovery_hash(self.hass, self._discovery_hash)
                 self._remove_signal()
             elif self._discovery_update:
                 # Non-empty payload: Notify component

homeassistant/components/mqtt/discovery.py

@@ -180,6 +180,11 @@ ABBREVIATIONS = {
 }
 
 
+def clear_discovery_hash(hass, discovery_hash):
+    """Clear entry in ALREADY_DISCOVERED list."""
+    del hass.data[ALREADY_DISCOVERED][discovery_hash]
+
+
 async def async_start(hass: HomeAssistantType, discovery_topic, hass_config,
                       config_entry=None) -> bool:
     """Initialize of MQTT Discovery."""

homeassistant/components/switch/mqtt.py

@@ -14,7 +14,8 @@ from homeassistant.components.mqtt import (
     CONF_AVAILABILITY_TOPIC, CONF_PAYLOAD_AVAILABLE,
     CONF_PAYLOAD_NOT_AVAILABLE, CONF_QOS, CONF_RETAIN, MqttAvailability,
     MqttDiscoveryUpdate, MqttEntityDeviceInfo, subscription)
-from homeassistant.components.mqtt.discovery import MQTT_DISCOVERY_NEW
+from homeassistant.components.mqtt.discovery import (
+    MQTT_DISCOVERY_NEW, clear_discovery_hash)
 from homeassistant.components.switch import SwitchDevice
 from homeassistant.const import (
     CONF_NAME, CONF_OPTIMISTIC, CONF_VALUE_TEMPLATE, CONF_PAYLOAD_OFF,
@@ -61,9 +62,15 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up MQTT switch dynamically through MQTT discovery."""
     async def async_discover(discovery_payload):
         """Discover and add a MQTT switch."""
+        try:
+            discovery_hash = discovery_payload[ATTR_DISCOVERY_HASH]
             config = PLATFORM_SCHEMA(discovery_payload)
             await _async_setup_entity(config, async_add_entities,
-                                  discovery_payload[ATTR_DISCOVERY_HASH])
+                                      discovery_hash)
+        except Exception:
+            if discovery_hash:
+                clear_discovery_hash(hass, discovery_hash)
+            raise
 
     async_dispatcher_connect(
         hass, MQTT_DISCOVERY_NEW.format(switch.DOMAIN, 'mqtt'),

tests/components/switch/test_mqtt.py

@@ -285,7 +285,7 @@ async def test_unique_id(hass):
 
 
 async def test_discovery_removal_switch(hass, mqtt_mock, caplog):
-    """Test expansion of discovered switch."""
+    """Test removal of discovered switch."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)
     await async_start(hass, 'homeassistant', {}, entry)
 
@@ -314,7 +314,7 @@ async def test_discovery_removal_switch(hass, mqtt_mock, caplog):
 
 
 async def test_discovery_update_switch(hass, mqtt_mock, caplog):
-    """Test expansion of discovered switch."""
+    """Test update of discovered switch."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)
     await async_start(hass, 'homeassistant', {}, entry)
 
@@ -349,6 +349,39 @@ async def test_discovery_update_switch(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/switch/bla/config',
+                            data1)
+    await hass.async_block_till_done()
+
+    state = hass.states.get('switch.beer')
+    assert state is None
+
+    async_fire_mqtt_message(hass, 'homeassistant/switch/bla/config',
+                            data2)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('switch.milk')
+    assert state is not None
+    assert state.name == 'Milk'
+    state = hass.states.get('switch.beer')
+    assert state is None
+
+
 async def test_entity_device_info_with_identifier(hass, mqtt_mock):
     """Test MQTT switch device registry integration."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN)