Cleanup Xiaomi Aqara (#10302)

homeassistant/components/xiaomi_aqara.py

@@ -1,4 +1,5 @@
 """Support for Xiaomi Gateways."""
+import asyncio
 import logging
 import voluptuous as vol
 import homeassistant.helpers.config_validation as cv
@@ -17,6 +18,7 @@ ATTR_DEVICE_ID = 'device_id'
 CONF_DISCOVERY_RETRY = 'discovery_retry'
 CONF_GATEWAYS = 'gateways'
 CONF_INTERFACE = 'interface'
+CONF_KEY = 'key'
 DOMAIN = 'xiaomi_aqara'
 PY_XIAOMI_GATEWAY = "xiaomi_gw"
 
@@ -25,76 +27,57 @@ SERVICE_STOP_RINGTONE = 'stop_ringtone'
 SERVICE_ADD_DEVICE = 'add_device'
 SERVICE_REMOVE_DEVICE = 'remove_device'
 
-XIAOMI_AQARA_SERVICE_SCHEMA = vol.Schema({
+
-    vol.Required(ATTR_GW_MAC): vol.All(cv.string,
+GW_MAC = vol.All(
-                                       vol.Any(vol.Length(min=12, max=12),
+    cv.string,
-                                               vol.Length(min=17, max=17)))
+    lambda value: value.replace(':', '').lower(),
+    vol.Length(min=12, max=12)
+)
+
+
+SERVICE_SCHEMA_PLAY_RINGTONE = vol.Schema({
+    vol.Required(ATTR_RINGTONE_ID):
+        vol.All(vol.Coerce(int), vol.NotIn([9, 14, 15, 16, 17, 18, 19])),
+    vol.Optional(ATTR_RINGTONE_VOL):
+        vol.All(vol.Coerce(int), vol.Clamp(min=0, max=100))
 })
 
-SERVICE_SCHEMA_PLAY_RINGTONE = XIAOMI_AQARA_SERVICE_SCHEMA.extend({
+SERVICE_SCHEMA_REMOVE_DEVICE = vol.Schema({
-    vol.Required(ATTR_RINGTONE_ID): vol.Coerce(int),
+    vol.Required(ATTR_DEVICE_ID):
-    vol.Optional(ATTR_RINGTONE_VOL): vol.All(vol.Coerce(int),
+        vol.All(cv.string, vol.Length(min=14, max=14))
-                                             vol.Clamp(min=0, max=100))
 })
 
-SERVICE_SCHEMA_REMOVE_DEVICE = XIAOMI_AQARA_SERVICE_SCHEMA.extend({
+
-    vol.Required(ATTR_DEVICE_ID): vol.All(cv.string,
+GATEWAY_CONFIG = vol.Schema({
-                                          vol.Length(min=14, max=14))
+    vol.Optional(CONF_MAC): GW_MAC,
+    vol.Optional(CONF_KEY, default=None):
+        vol.All(cv.string, vol.Length(min=16, max=16)),
+    vol.Optional(CONF_HOST): cv.string,
+    vol.Optional(CONF_PORT, default=9898): cv.port,
 })
 
-SERVICE_TO_METHOD = {
+
-    SERVICE_PLAY_RINGTONE: {'method': 'play_ringtone_service',
+def _fix_conf_defaults(config):
-                            'schema': SERVICE_SCHEMA_PLAY_RINGTONE},
+    """Update some config defaults."""
-    SERVICE_STOP_RINGTONE: {'method': 'stop_ringtone_service'},
+    config['sid'] = config.pop(CONF_MAC, None)
-    SERVICE_ADD_DEVICE: {'method': 'add_device_service'},
+
-    SERVICE_REMOVE_DEVICE: {'method': 'remove_device_service',
+    if config.get(CONF_KEY) is None:
-                            'schema': SERVICE_SCHEMA_REMOVE_DEVICE},
+        _LOGGER.warning(
-}
+            'Key is not provided for gateway %s. Controlling the gateway '
+            'will not be possible.', config['sid'])
+
+    if config.get(CONF_HOST) is None:
+        config.pop(CONF_PORT)
+
+    return config
 
 
-def _validate_conf(config):
+DEFAULT_GATEWAY_CONFIG = [{CONF_MAC: None, CONF_KEY: None}]
-    """Validate a list of devices definitions."""
-    res_config = []
-    for gw_conf in config:
-        for _conf in gw_conf.keys():
-            if _conf not in [CONF_MAC, CONF_HOST, CONF_PORT, 'key']:
-                raise vol.Invalid('{} is not a valid config parameter'.
-                                  format(_conf))
-
-        res_gw_conf = {'sid': gw_conf.get(CONF_MAC)}
-        if res_gw_conf['sid'] is not None:
-            res_gw_conf['sid'] = res_gw_conf['sid'].replace(":", "").lower()
-            if len(res_gw_conf['sid']) != 12:
-                raise vol.Invalid('Invalid mac address', gw_conf.get(CONF_MAC))
-        key = gw_conf.get('key')
-
-        if key is None:
-            _LOGGER.warning(
-                'Gateway Key is not provided.'
-                ' Controlling gateway device will not be possible.')
-        elif len(key) != 16:
-            raise vol.Invalid('Invalid key {}.'
-                              ' Key must be 16 characters'.format(key))
-        res_gw_conf['key'] = key
-
-        host = gw_conf.get(CONF_HOST)
-        if host is not None:
-            res_gw_conf[CONF_HOST] = host
-            res_gw_conf['port'] = gw_conf.get(CONF_PORT, 9898)
-
-            _LOGGER.warning(
-                'Static address (%s:%s) of the gateway provided. '
-                'Discovery of this host will be skipped.',
-                res_gw_conf[CONF_HOST], res_gw_conf[CONF_PORT])
-
-        res_config.append(res_gw_conf)
-    return res_config
-
 
 CONFIG_SCHEMA = vol.Schema({
     DOMAIN: vol.Schema({
-        vol.Optional(CONF_GATEWAYS, default=[{CONF_MAC: None, "key": None}]):
+        vol.Optional(CONF_GATEWAYS, default=DEFAULT_GATEWAY_CONFIG):
-            vol.All(cv.ensure_list, _validate_conf),
+            vol.All(cv.ensure_list, [GATEWAY_CONFIG], [_fix_conf_defaults]),
         vol.Optional(CONF_INTERFACE, default='any'): cv.string,
         vol.Optional(CONF_DISCOVERY_RETRY, default=3): cv.positive_int
     })
@@ -113,30 +96,30 @@ def setup(hass, config):
         interface = config[DOMAIN][CONF_INTERFACE]
         discovery_retry = config[DOMAIN][CONF_DISCOVERY_RETRY]
 
+    @asyncio.coroutine
     def xiaomi_gw_discovered(service, discovery_info):
         """Called when Xiaomi Gateway device(s) has been found."""
         # We don't need to do anything here, the purpose of HA's
         # discovery service is to just trigger loading of this
         # component, and then its own discovery process kicks in.
-        _LOGGER.info("Discovered: %s", discovery_info)
 
     discovery.listen(hass, SERVICE_XIAOMI_GW, xiaomi_gw_discovered)
 
     from PyXiaomiGateway import PyXiaomiGateway
-    hass.data[PY_XIAOMI_GATEWAY] = PyXiaomiGateway(hass.add_job, gateways,
+    xiaomi = hass.data[PY_XIAOMI_GATEWAY] = PyXiaomiGateway(
-                                                   interface)
+        hass.add_job, gateways, interface)
 
     _LOGGER.debug("Expecting %s gateways", len(gateways))
     for k in range(discovery_retry):
         _LOGGER.info('Discovering Xiaomi Gateways (Try %s)', k + 1)
-        hass.data[PY_XIAOMI_GATEWAY].discover_gateways()
+        xiaomi.discover_gateways()
-        if len(hass.data[PY_XIAOMI_GATEWAY].gateways) >= len(gateways):
+        if len(xiaomi.gateways) >= len(gateways):
             break
 
-    if not hass.data[PY_XIAOMI_GATEWAY].gateways:
+    if not xiaomi.gateways:
         _LOGGER.error("No gateway discovered")
         return False
-    hass.data[PY_XIAOMI_GATEWAY].listen()
+    xiaomi.listen()
     _LOGGER.debug("Gateways discovered. Listening for broadcasts")
 
     for component in ['binary_sensor', 'sensor', 'switch', 'light', 'cover']:
@@ -145,81 +128,60 @@ def setup(hass, config):
     def stop_xiaomi(event):
         """Stop Xiaomi Socket."""
         _LOGGER.info("Shutting down Xiaomi Hub.")
-        hass.data[PY_XIAOMI_GATEWAY].stop_listen()
+        xiaomi.stop_listen()
+
     hass.bus.listen_once(EVENT_HOMEASSISTANT_STOP, stop_xiaomi)
 
-    # pylint: disable=unused-variable
     def play_ringtone_service(call):
         """Service to play ringtone through Gateway."""
-        ring_id = int(call.data.get(ATTR_RINGTONE_ID))
+        ring_id = call.data.get(ATTR_RINGTONE_ID)
-        gw_sid = call.data.get(ATTR_GW_MAC).replace(":", "").lower()
+        gateway = call.data.get(ATTR_GW_MAC)
 
-        if ring_id in [9, 14-19]:
+        kwargs = {'mid': ring_id}
-            _LOGGER.error('Specified mid: %s is not defined in gateway.',
-                          ring_id)
-            return
 
         ring_vol = call.data.get(ATTR_RINGTONE_VOL)
-        if ring_vol is None:
+        if ring_vol is not None:
-            ringtone = {'mid': ring_id}
+            kwargs['vol'] = ring_vol
-        else:
-            ringtone = {'mid': ring_id, 'vol': int(ring_vol)}
 
-        for (_, gateway) in hass.data[PY_XIAOMI_GATEWAY].gateways.items():
+        gateway.write_to_hub(gateway.sid, **kwargs)
-            if gateway.sid == gw_sid:
-                gateway.write_to_hub(gateway.sid, **ringtone)
-                break
-        else:
-            _LOGGER.error('Unknown gateway sid: %s was specified.', gw_sid)
 
-    # pylint: disable=unused-variable
     def stop_ringtone_service(call):
         """Service to stop playing ringtone on Gateway."""
-        gw_sid = call.data.get(ATTR_GW_MAC).replace(":", "").lower()
+        gateway = call.data.get(ATTR_GW_MAC)
-        for (_, gateway) in hass.data[PY_XIAOMI_GATEWAY].gateways.items():
+        gateway.write_to_hub(gateway.sid, mid=10000)
-            if gateway.sid == gw_sid:
-                ringtone = {'mid': 10000}
-                gateway.write_to_hub(gateway.sid, **ringtone)
-                break
-        else:
-            _LOGGER.error('Unknown gateway sid: %s was specified.', gw_sid)
 
-    # pylint: disable=unused-variable
     def add_device_service(call):
         """Service to add a new sub-device within the next 30 seconds."""
-        gw_sid = call.data.get(ATTR_GW_MAC).replace(":", "").lower()
+        gateway = call.data.get(ATTR_GW_MAC)
-        for (_, gateway) in hass.data[PY_XIAOMI_GATEWAY].gateways.items():
+        gateway.write_to_hub(gateway.sid, join_permission='yes')
-            if gateway.sid == gw_sid:
+        hass.components.persistent_notification.async_create(
-                join_permission = {'join_permission': 'yes'}
+            'Join permission enabled for 30 seconds! '
-                gateway.write_to_hub(gateway.sid, **join_permission)
+            'Please press the pairing button of the new device once.',
-                hass.components.persistent_notification.async_create(
+            title='Xiaomi Aqara Gateway')
-                    'Join permission enabled for 30 seconds! '
-                    'Please press the pairing button of the new device once.',
-                    title='Xiaomi Aqara Gateway')
-                break
-        else:
-            _LOGGER.error('Unknown gateway sid: %s was specified.', gw_sid)
 
-    # pylint: disable=unused-variable
     def remove_device_service(call):
         """Service to remove a sub-device from the gateway."""
         device_id = call.data.get(ATTR_DEVICE_ID)
-        gw_sid = call.data.get(ATTR_GW_MAC).replace(":", "").lower()
+        gateway = call.data.get(ATTR_GW_MAC)
-        remove_device = {'remove_device': device_id}
+        gateway.write_to_hub(gateway.sid, remove_device=device_id)
-        for (_, gateway) in hass.data[PY_XIAOMI_GATEWAY].gateways.items():
-            if gateway.sid == gw_sid:
-                gateway.write_to_hub(gateway.sid, **remove_device)
-                break
-        else:
-            _LOGGER.error('Unknown gateway sid: %s was specified.', gw_sid)
 
-    for xiaomi_aqara_service in SERVICE_TO_METHOD:
+    gateway_only_schema = _add_gateway_to_schema(xiaomi, vol.Schema({}))
-        schema = SERVICE_TO_METHOD[xiaomi_aqara_service].get(
+
-            'schema', XIAOMI_AQARA_SERVICE_SCHEMA)
+    hass.services.async_register(
-        service_handler = SERVICE_TO_METHOD[xiaomi_aqara_service].get('method')
+        DOMAIN, SERVICE_PLAY_RINGTONE, play_ringtone_service,
-        hass.services.async_register(
+        schema=_add_gateway_to_schema(xiaomi, SERVICE_SCHEMA_PLAY_RINGTONE))
-            DOMAIN, xiaomi_aqara_service, service_handler,
+
-            description=None, schema=schema)
+    hass.services.async_register(
+        DOMAIN, SERVICE_STOP_RINGTONE, stop_ringtone_service,
+        schema=gateway_only_schema)
+
+    hass.services.async_register(
+        DOMAIN, SERVICE_ADD_DEVICE, add_device_service,
+        schema=gateway_only_schema)
+
+    hass.services.async_register(
+        DOMAIN, SERVICE_REMOVE_DEVICE, remove_device_service,
+        schema=_add_gateway_to_schema(xiaomi, SERVICE_SCHEMA_REMOVE_DEVICE))
 
     return True
 
@@ -276,3 +238,27 @@ class XiaomiDevice(Entity):
     def parse_data(self, data):
         """Parse data sent by gateway."""
         raise NotImplementedError()
+
+
+def _add_gateway_to_schema(xiaomi, schema):
+    """Extend a voluptuous schema with a gateway validator."""
+    def gateway(sid):
+        """Convert sid to a gateway."""
+        sid = str(sid).replace(':', '').lower()
+
+        for gateway in xiaomi.gateways.values():
+            if gateway.sid == sid:
+                return gateway
+
+        raise vol.Invalid('Unknown gateway sid {}'.format(sid))
+
+    gateways = list(xiaomi.gateways.values())
+    kwargs = {}
+
+    # If the user has only 1 gateway, make it the default for services.
+    if len(gateways) == 1:
+        kwargs['default'] = gateways[0]
+
+    return schema.extend({
+        vol.Required(ATTR_GW_MAC, **kwargs): gateway
+    })