Update device registry of MQTT Switch (#19540)

* MQTT Switch: Update device registry

* Move config_entry to constructor

* Remove duplicated code

* Fix merge error
This commit is contained in:
emontnemery 2019-01-25 14:39:16 +08:00 committed by Paulus Schoutsen
parent 55943cfac0
commit a1da6a677a
3 changed files with 70 additions and 8 deletions

View file

@ -65,7 +65,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
try:
discovery_hash = discovery_payload[ATTR_DISCOVERY_HASH]
config = PLATFORM_SCHEMA(discovery_payload)
await _async_setup_entity(config, async_add_entities,
await _async_setup_entity(config, async_add_entities, config_entry,
discovery_hash)
except Exception:
if discovery_hash:
@ -77,10 +77,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
async_discover)
async def _async_setup_entity(config, async_add_entities,
async def _async_setup_entity(config, async_add_entities, config_entry=None,
discovery_hash=None):
"""Set up the MQTT switch."""
async_add_entities([MqttSwitch(config, discovery_hash)])
async_add_entities([MqttSwitch(config, config_entry, discovery_hash)])
# pylint: disable=too-many-ancestors
@ -88,7 +88,7 @@ class MqttSwitch(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
MqttEntityDeviceInfo, SwitchDevice, RestoreEntity):
"""Representation of a switch that can be toggled using MQTT."""
def __init__(self, config, discovery_hash):
def __init__(self, config, config_entry, discovery_hash):
"""Initialize the MQTT switch."""
self._state = False
self._sub_state = None
@ -107,7 +107,7 @@ class MqttSwitch(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
MqttAvailability.__init__(self, config)
MqttDiscoveryUpdate.__init__(self, discovery_hash,
self.discovery_update)
MqttEntityDeviceInfo.__init__(self, device_config)
MqttEntityDeviceInfo.__init__(self, device_config, config_entry)
async def async_added_to_hass(self):
"""Subscribe to MQTT events."""
@ -120,6 +120,7 @@ class MqttSwitch(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
self._setup_from_config(config)
await self.attributes_discovery_update(config)
await self.availability_discovery_update(config)
await self.device_info_discovery_update(config)
await self._subscribe_topics()
self.async_schedule_update_ha_state()