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
This commit is contained in:
emontnemery 2019-01-25 14:40:52 +08:00 committed by Paulus Schoutsen
parent a1da6a677a
commit d84cd01cbf
7 changed files with 132 additions and 16 deletions

View File

@ -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)

View File

@ -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(

View File

@ -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)])

View File

@ -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)])

View File

@ -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)

View File

@ -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)

View File

@ -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)