Refactor ZHA gateway into modules and add admin protections to API (#22023)

* refactor

* cleanup

* fix tests

* admin all the things

homeassistant/components/zha/__init__.py

@@ -4,10 +4,7 @@ Support for Zigbee Home Automation devices.
 For more details about this component, please refer to the documentation at
 https://home-assistant.io/components/zha/
 """
-import asyncio
 import logging
-import os
-import types
 
 import voluptuous as vol
 
@@ -21,13 +18,13 @@ from . import api
 from .core import ZHAGateway
 from .core.const import (
     COMPONENTS, CONF_BAUDRATE, CONF_DATABASE, CONF_DEVICE_CONFIG,
-    CONF_RADIO_TYPE, CONF_USB_PATH, DATA_ZHA, DATA_ZHA_BRIDGE_ID,
+    CONF_RADIO_TYPE, CONF_USB_PATH, DATA_ZHA,
     DATA_ZHA_CONFIG, DATA_ZHA_CORE_COMPONENT, DATA_ZHA_DISPATCHERS,
-    DATA_ZHA_RADIO, DEFAULT_BAUDRATE, DEFAULT_DATABASE_NAME, DATA_ZHA_GATEWAY,
+    DATA_ZHA_RADIO, DEFAULT_BAUDRATE, DATA_ZHA_GATEWAY,
     DEFAULT_RADIO_TYPE, DOMAIN, RadioType, DATA_ZHA_CORE_EVENTS, ENABLE_QUIRKS)
-from .core.gateway import establish_device_mappings
+from .core.registries import establish_device_mappings
 from .core.channels.registry import populate_channel_registry
-from .core.store import async_get_registry
+from .core.patches import apply_cluster_listener_patch
 
 REQUIREMENTS = [
     'bellows-homeassistant==0.7.1',
@@ -108,82 +105,32 @@ async def async_setup_entry(hass, config_entry):
         # pylint: disable=W0611, W0612
         import zhaquirks  # noqa
 
-    usb_path = config_entry.data.get(CONF_USB_PATH)
-    baudrate = config.get(CONF_BAUDRATE, DEFAULT_BAUDRATE)
-    radio_type = config_entry.data.get(CONF_RADIO_TYPE)
-    if radio_type == RadioType.ezsp.name:
-        import bellows.ezsp
-        from bellows.zigbee.application import ControllerApplication
-        radio = bellows.ezsp.EZSP()
-        radio_description = "EZSP"
-    elif radio_type == RadioType.xbee.name:
-        import zigpy_xbee.api
-        from zigpy_xbee.zigbee.application import ControllerApplication
-        radio = zigpy_xbee.api.XBee()
-        radio_description = "XBee"
-    elif radio_type == RadioType.deconz.name:
-        import zigpy_deconz.api
-        from zigpy_deconz.zigbee.application import ControllerApplication
-        radio = zigpy_deconz.api.Deconz()
-        radio_description = "Deconz"
-
-    await radio.connect(usb_path, baudrate)
-    hass.data[DATA_ZHA][DATA_ZHA_RADIO] = radio
-
-    if CONF_DATABASE in config:
-        database = config[CONF_DATABASE]
-    else:
-        database = os.path.join(hass.config.config_dir, DEFAULT_DATABASE_NAME)
-
     # patch zigpy listener to prevent flooding logs with warnings due to
     # how zigpy implemented its listeners
-    from zigpy.appdb import ClusterPersistingListener
+    apply_cluster_listener_patch()
 
-    def zha_send_event(self, cluster, command, args):
+    zha_gateway = ZHAGateway(hass, config)
-        pass
+    await zha_gateway.async_initialize(config_entry)
-
-    ClusterPersistingListener.zha_send_event = types.MethodType(
-        zha_send_event,
-        ClusterPersistingListener
-    )
-
-    zha_storage = await async_get_registry(hass)
-    zha_gateway = ZHAGateway(hass, config, zha_storage)
-
-    # Patch handle_message until zigpy can provide an event here
-    def handle_message(sender, is_reply, profile, cluster,
-                       src_ep, dst_ep, tsn, command_id, args):
-        """Handle message from a device."""
-        if not sender.initializing and sender.ieee in zha_gateway.devices and \
-                not zha_gateway.devices[sender.ieee].available:
-            zha_gateway.async_device_became_available(
-                sender, is_reply, profile, cluster, src_ep, dst_ep, tsn,
-                command_id, args
-            )
-        return sender.handle_message(
-            is_reply, profile, cluster, src_ep, dst_ep, tsn, command_id, args)
-
-    application_controller = ControllerApplication(radio, database)
-    application_controller.handle_message = handle_message
-    application_controller.add_listener(zha_gateway)
-    await application_controller.startup(auto_form=True)
-
-    hass.data[DATA_ZHA][DATA_ZHA_BRIDGE_ID] = str(application_controller.ieee)
-
-    init_tasks = []
-    for device in application_controller.devices.values():
-        init_tasks.append(zha_gateway.async_device_initialized(device, False))
-    await asyncio.gather(*init_tasks)
 
     device_registry = await \
         hass.helpers.device_registry.async_get_registry()
     device_registry.async_get_or_create(
         config_entry_id=config_entry.entry_id,
-        connections={(CONNECTION_ZIGBEE, str(application_controller.ieee))},
+        connections={
-        identifiers={(DOMAIN, str(application_controller.ieee))},
+            (
+                CONNECTION_ZIGBEE,
+                str(zha_gateway.application_controller.ieee)
+            )
+        },
+        identifiers={
+            (
+                DOMAIN,
+                str(zha_gateway.application_controller.ieee)
+            )
+        },
         name="Zigbee Coordinator",
         manufacturer="ZHA",
-        model=radio_description,
+        model=zha_gateway.radio_description,
     )
 
     for component in COMPONENTS:
@@ -192,7 +139,7 @@ async def async_setup_entry(hass, config_entry):
                 config_entry, component)
         )
 
-    api.async_load_api(hass, application_controller, zha_gateway)
+    api.async_load_api(hass)
 
     async def async_zha_shutdown(event):
         """Handle shutdown tasks."""

homeassistant/components/zha/api.py

@@ -74,6 +74,7 @@ SERVICE_SCHEMAS = {
 }
 
 
+@websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command({
     vol.Required(TYPE): 'zha/devices'
@@ -103,6 +104,7 @@ async def websocket_get_devices(hass, connection, msg):
     connection.send_result(msg[ID], devices)
 
 
+@websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command({
     vol.Required(TYPE): 'zha/devices/reconfigure',
@@ -117,6 +119,7 @@ async def websocket_reconfigure_node(hass, connection, msg):
     hass.async_create_task(device.async_configure())
 
 
+@websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command({
     vol.Required(TYPE): 'zha/devices/clusters',
@@ -149,6 +152,7 @@ async def websocket_device_clusters(hass, connection, msg):
     connection.send_result(msg[ID], response_clusters)
 
 
+@websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command({
     vol.Required(TYPE): 'zha/devices/clusters/attributes',
@@ -190,6 +194,7 @@ async def websocket_device_cluster_attributes(hass, connection, msg):
     connection.send_result(msg[ID], cluster_attributes)
 
 
+@websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command({
     vol.Required(TYPE): 'zha/devices/clusters/commands',
@@ -241,6 +246,7 @@ async def websocket_device_cluster_commands(hass, connection, msg):
     connection.send_result(msg[ID], cluster_commands)
 
 
+@websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command({
     vol.Required(TYPE): 'zha/devices/clusters/attributes/value',
@@ -283,6 +289,7 @@ async def websocket_read_zigbee_cluster_attributes(hass, connection, msg):
     connection.send_result(msg[ID], str(success.get(attribute)))
 
 
+@websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command({
     vol.Required(TYPE): 'zha/devices/bindable',
@@ -311,6 +318,7 @@ async def websocket_get_bindable_devices(hass, connection, msg):
     ))
 
 
+@websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command({
     vol.Required(TYPE): 'zha/devices/bind',
@@ -330,6 +338,7 @@ async def websocket_bind_devices(hass, connection, msg):
                  )
 
 
+@websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command({
     vol.Required(TYPE): 'zha/devices/unbind',
@@ -386,16 +395,19 @@ async def async_binding_operation(zha_gateway, source_ieee, target_ieee,
     await asyncio.gather(*bind_tasks)
 
 
-def async_load_api(hass, application_controller, zha_gateway):
+def async_load_api(hass):
     """Set up the web socket API."""
+    zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
+    application_controller = zha_gateway.application_controller
+
     async def permit(service):
         """Allow devices to join this network."""
         duration = service.data.get(ATTR_DURATION)
         _LOGGER.info("Permitting joins for %ss", duration)
         await application_controller.permit(duration)
 
-    hass.services.async_register(DOMAIN, SERVICE_PERMIT, permit,
+    hass.helpers.service.async_register_admin_service(
-                                 schema=SERVICE_SCHEMAS[SERVICE_PERMIT])
+        DOMAIN, SERVICE_PERMIT, permit, schema=SERVICE_SCHEMAS[SERVICE_PERMIT])
 
     async def remove(service):
         """Remove a node from the network."""
@@ -405,8 +417,8 @@ def async_load_api(hass, application_controller, zha_gateway):
         _LOGGER.info("Removing node %s", ieee)
         await application_controller.remove(ieee)
 
-    hass.services.async_register(DOMAIN, SERVICE_REMOVE, remove,
+    hass.helpers.service.async_register_admin_service(
-                                 schema=SERVICE_SCHEMAS[IEEE_SERVICE])
+        DOMAIN, SERVICE_REMOVE, remove, schema=SERVICE_SCHEMAS[IEEE_SERVICE])
 
     async def set_zigbee_cluster_attributes(service):
         """Set zigbee attribute for cluster on zha entity."""
@@ -438,11 +450,12 @@ def async_load_api(hass, application_controller, zha_gateway):
                       "{}: [{}]".format(RESPONSE, response)
                       )
 
-    hass.services.async_register(DOMAIN, SERVICE_SET_ZIGBEE_CLUSTER_ATTRIBUTE,
+    hass.helpers.service.async_register_admin_service(
-                                 set_zigbee_cluster_attributes,
+        DOMAIN, SERVICE_SET_ZIGBEE_CLUSTER_ATTRIBUTE,
-                                 schema=SERVICE_SCHEMAS[
+        set_zigbee_cluster_attributes,
-                                     SERVICE_SET_ZIGBEE_CLUSTER_ATTRIBUTE
+        schema=SERVICE_SCHEMAS[
-                                 ])
+            SERVICE_SET_ZIGBEE_CLUSTER_ATTRIBUTE
+        ])
 
     async def issue_zigbee_cluster_command(service):
         """Issue command on zigbee cluster on zha entity."""
@@ -477,11 +490,12 @@ def async_load_api(hass, application_controller, zha_gateway):
                       "{}: [{}]".format(RESPONSE, response)
                       )
 
-    hass.services.async_register(DOMAIN, SERVICE_ISSUE_ZIGBEE_CLUSTER_COMMAND,
+    hass.helpers.service.async_register_admin_service(
-                                 issue_zigbee_cluster_command,
+        DOMAIN, SERVICE_ISSUE_ZIGBEE_CLUSTER_COMMAND,
-                                 schema=SERVICE_SCHEMAS[
+        issue_zigbee_cluster_command,
-                                     SERVICE_ISSUE_ZIGBEE_CLUSTER_COMMAND
+        schema=SERVICE_SCHEMAS[
-                                 ])
+            SERVICE_ISSUE_ZIGBEE_CLUSTER_COMMAND
+        ])
 
     websocket_api.async_register_command(hass, websocket_get_devices)
     websocket_api.async_register_command(hass, websocket_reconfigure_node)

homeassistant/components/zha/core/const.py

@@ -33,6 +33,10 @@ CONF_USB_PATH = 'usb_path'
 DATA_DEVICE_CONFIG = 'zha_device_config'
 ENABLE_QUIRKS = 'enable_quirks'
 
+RADIO = 'radio'
+RADIO_DESCRIPTION = 'radio_description'
+CONTROLLER = 'controller'
+
 DEFAULT_RADIO_TYPE = 'ezsp'
 DEFAULT_BAUDRATE = 57600
 DEFAULT_DATABASE_NAME = 'zigbee.db'
@@ -114,6 +118,9 @@ DISCOVERY_KEY = 'zha_discovery_info'
 DEVICE_CLASS = {}
 SINGLE_INPUT_CLUSTER_DEVICE_CLASS = {}
 SINGLE_OUTPUT_CLUSTER_DEVICE_CLASS = {}
+SENSOR_TYPES = {}
+RADIO_TYPES = {}
+BINARY_SENSOR_TYPES = {}
 CLUSTER_REPORT_CONFIGS = {}
 CUSTOM_CLUSTER_MAPPINGS = {}
 COMPONENT_CLUSTERS = {}

homeassistant/components/zha/core/discovery.py

@@ -0,0 +1,265 @@
+"""
+Device discovery functions for Zigbee Home Automation.
+
+For more details about this component, please refer to the documentation at
+https://home-assistant.io/components/zha/
+"""
+
+import logging
+
+from homeassistant import const as ha_const
+from homeassistant.core import callback
+from homeassistant.helpers.dispatcher import async_dispatcher_send
+from . import const as zha_const
+from .channels import (
+    AttributeListeningChannel, EventRelayChannel, ZDOChannel
+)
+from .channels.registry import ZIGBEE_CHANNEL_REGISTRY
+from .const import (
+    CONF_DEVICE_CONFIG, COMPONENTS, ZHA_DISCOVERY_NEW, DATA_ZHA,
+    SENSOR_TYPE, UNKNOWN, BINARY_SENSOR_TYPES, NO_SENSOR_CLUSTERS,
+    EVENT_RELAY_CLUSTERS, SENSOR_TYPES, GENERIC,
+    POWER_CONFIGURATION_CHANNEL
+)
+from ..device_entity import ZhaDeviceEntity
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@callback
+def async_process_endpoint(
+        hass, config, endpoint_id, endpoint, discovery_infos, device,
+        zha_device, is_new_join):
+    """Process an endpoint on a zigpy device."""
+    import zigpy.profiles
+
+    if endpoint_id == 0:  # ZDO
+        _async_create_cluster_channel(
+            endpoint,
+            zha_device,
+            is_new_join,
+            channel_class=ZDOChannel
+        )
+        return
+
+    component = None
+    profile_clusters = ([], [])
+    device_key = "{}-{}".format(device.ieee, endpoint_id)
+    node_config = {}
+    if CONF_DEVICE_CONFIG in config:
+        node_config = config[CONF_DEVICE_CONFIG].get(
+            device_key, {}
+        )
+
+    if endpoint.profile_id in zigpy.profiles.PROFILES:
+        profile = zigpy.profiles.PROFILES[endpoint.profile_id]
+        if zha_const.DEVICE_CLASS.get(endpoint.profile_id,
+                                      {}).get(endpoint.device_type,
+                                              None):
+            profile_clusters = profile.CLUSTERS[endpoint.device_type]
+            profile_info = zha_const.DEVICE_CLASS[endpoint.profile_id]
+            component = profile_info[endpoint.device_type]
+
+    if ha_const.CONF_TYPE in node_config:
+        component = node_config[ha_const.CONF_TYPE]
+        profile_clusters = zha_const.COMPONENT_CLUSTERS[component]
+
+    if component and component in COMPONENTS:
+        profile_match = _async_handle_profile_match(
+            hass, endpoint, profile_clusters, zha_device,
+            component, device_key, is_new_join)
+        discovery_infos.append(profile_match)
+
+    discovery_infos.extend(_async_handle_single_cluster_matches(
+        hass,
+        endpoint,
+        zha_device,
+        profile_clusters,
+        device_key,
+        is_new_join
+    ))
+
+
+@callback
+def _async_create_cluster_channel(cluster, zha_device, is_new_join,
+                                  channels=None, channel_class=None):
+    """Create a cluster channel and attach it to a device."""
+    if channel_class is None:
+        channel_class = ZIGBEE_CHANNEL_REGISTRY.get(cluster.cluster_id,
+                                                    AttributeListeningChannel)
+    channel = channel_class(cluster, zha_device)
+    zha_device.add_cluster_channel(channel)
+    if channels is not None:
+        channels.append(channel)
+
+
+@callback
+def async_dispatch_discovery_info(hass, is_new_join, discovery_info):
+    """Dispatch or store discovery information."""
+    if not discovery_info['channels']:
+        _LOGGER.warning(
+            "there are no channels in the discovery info: %s", discovery_info)
+        return
+    component = discovery_info['component']
+    if is_new_join:
+        async_dispatcher_send(
+            hass,
+            ZHA_DISCOVERY_NEW.format(component),
+            discovery_info
+        )
+    else:
+        hass.data[DATA_ZHA][component][discovery_info['unique_id']] = \
+            discovery_info
+
+
+@callback
+def _async_handle_profile_match(hass, endpoint, profile_clusters, zha_device,
+                                component, device_key, is_new_join):
+    """Dispatch a profile match to the appropriate HA component."""
+    in_clusters = [endpoint.in_clusters[c]
+                   for c in profile_clusters[0]
+                   if c in endpoint.in_clusters]
+    out_clusters = [endpoint.out_clusters[c]
+                    for c in profile_clusters[1]
+                    if c in endpoint.out_clusters]
+
+    channels = []
+
+    for cluster in in_clusters:
+        _async_create_cluster_channel(
+            cluster, zha_device, is_new_join, channels=channels)
+
+    for cluster in out_clusters:
+        _async_create_cluster_channel(
+            cluster, zha_device, is_new_join, channels=channels)
+
+    discovery_info = {
+        'unique_id': device_key,
+        'zha_device': zha_device,
+        'channels': channels,
+        'component': component
+    }
+
+    if component == 'binary_sensor':
+        discovery_info.update({SENSOR_TYPE: UNKNOWN})
+        cluster_ids = []
+        cluster_ids.extend(profile_clusters[0])
+        cluster_ids.extend(profile_clusters[1])
+        for cluster_id in cluster_ids:
+            if cluster_id in BINARY_SENSOR_TYPES:
+                discovery_info.update({
+                    SENSOR_TYPE: BINARY_SENSOR_TYPES.get(
+                        cluster_id, UNKNOWN)
+                })
+                break
+
+    return discovery_info
+
+
+@callback
+def _async_handle_single_cluster_matches(hass, endpoint, zha_device,
+                                         profile_clusters, device_key,
+                                         is_new_join):
+    """Dispatch single cluster matches to HA components."""
+    cluster_matches = []
+    cluster_match_results = []
+    for cluster in endpoint.in_clusters.values():
+        # don't let profiles prevent these channels from being created
+        if cluster.cluster_id in NO_SENSOR_CLUSTERS:
+            cluster_match_results.append(
+                _async_handle_channel_only_cluster_match(
+                    zha_device,
+                    cluster,
+                    is_new_join,
+                ))
+
+        if cluster.cluster_id not in profile_clusters[0]:
+            cluster_match_results.append(_async_handle_single_cluster_match(
+                hass,
+                zha_device,
+                cluster,
+                device_key,
+                zha_const.SINGLE_INPUT_CLUSTER_DEVICE_CLASS,
+                is_new_join,
+            ))
+
+    for cluster in endpoint.out_clusters.values():
+        if cluster.cluster_id not in profile_clusters[1]:
+            cluster_match_results.append(_async_handle_single_cluster_match(
+                hass,
+                zha_device,
+                cluster,
+                device_key,
+                zha_const.SINGLE_OUTPUT_CLUSTER_DEVICE_CLASS,
+                is_new_join,
+            ))
+
+        if cluster.cluster_id in EVENT_RELAY_CLUSTERS:
+            _async_create_cluster_channel(
+                cluster,
+                zha_device,
+                is_new_join,
+                channel_class=EventRelayChannel
+            )
+
+    for cluster_match in cluster_match_results:
+        if cluster_match is not None:
+            cluster_matches.append(cluster_match)
+    return cluster_matches
+
+
+@callback
+def _async_handle_channel_only_cluster_match(
+        zha_device, cluster, is_new_join):
+    """Handle a channel only cluster match."""
+    _async_create_cluster_channel(cluster, zha_device, is_new_join)
+
+
+@callback
+def _async_handle_single_cluster_match(hass, zha_device, cluster, device_key,
+                                       device_classes, is_new_join):
+    """Dispatch a single cluster match to a HA component."""
+    component = None  # sub_component = None
+    for cluster_type, candidate_component in device_classes.items():
+        if isinstance(cluster_type, int):
+            if cluster.cluster_id == cluster_type:
+                component = candidate_component
+        elif isinstance(cluster, cluster_type):
+            component = candidate_component
+            break
+
+    if component is None or component not in COMPONENTS:
+        return
+    channels = []
+    _async_create_cluster_channel(cluster, zha_device, is_new_join,
+                                  channels=channels)
+
+    cluster_key = "{}-{}".format(device_key, cluster.cluster_id)
+    discovery_info = {
+        'unique_id': cluster_key,
+        'zha_device': zha_device,
+        'channels': channels,
+        'entity_suffix': '_{}'.format(cluster.cluster_id),
+        'component': component
+    }
+
+    if component == 'sensor':
+        discovery_info.update({
+            SENSOR_TYPE: SENSOR_TYPES.get(cluster.cluster_id, GENERIC)
+        })
+    if component == 'binary_sensor':
+        discovery_info.update({
+            SENSOR_TYPE: BINARY_SENSOR_TYPES.get(cluster.cluster_id, UNKNOWN)
+        })
+
+    return discovery_info
+
+
+@callback
+def async_create_device_entity(zha_device):
+    """Create ZHADeviceEntity."""
+    device_entity_channels = []
+    if POWER_CONFIGURATION_CHANNEL in zha_device.cluster_channels:
+        channel = zha_device.cluster_channels.get(POWER_CONFIGURATION_CHANNEL)
+        device_entity_channels.append(channel)
+    return ZhaDeviceEntity(zha_device, device_entity_channels)

homeassistant/components/zha/core/gateway.py

@@ -5,39 +5,35 @@ For more details about this component, please refer to the documentation at
 https://home-assistant.io/components/zha/
 """
 
+import asyncio
 import collections
 import itertools
 import logging
-from homeassistant import const as ha_const
+import os
+
 from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.entity_component import EntityComponent
-from . import const as zha_const
 from .const import (
-    COMPONENTS, CONF_DEVICE_CONFIG, DATA_ZHA, DATA_ZHA_CORE_COMPONENT, DOMAIN,
+    DATA_ZHA, DATA_ZHA_CORE_COMPONENT, DOMAIN,
-    ZHA_DISCOVERY_NEW, DEVICE_CLASS, SINGLE_INPUT_CLUSTER_DEVICE_CLASS,
+    SIGNAL_REMOVE, DATA_ZHA_GATEWAY, CONF_USB_PATH, CONF_BAUDRATE,
-    SINGLE_OUTPUT_CLUSTER_DEVICE_CLASS, COMPONENT_CLUSTERS, HUMIDITY,
+    DEFAULT_BAUDRATE, CONF_RADIO_TYPE, DATA_ZHA_RADIO, CONF_DATABASE,
-    TEMPERATURE, ILLUMINANCE, PRESSURE, METERING, ELECTRICAL_MEASUREMENT,
+    DEFAULT_DATABASE_NAME, DATA_ZHA_BRIDGE_ID, RADIO_TYPES,
-    GENERIC, SENSOR_TYPE, EVENT_RELAY_CLUSTERS, UNKNOWN, OPENING, ZONE,
+    RADIO, CONTROLLER, RADIO_DESCRIPTION)
-    OCCUPANCY, CLUSTER_REPORT_CONFIGS, REPORT_CONFIG_IMMEDIATE,
-    REPORT_CONFIG_ASAP, REPORT_CONFIG_DEFAULT, REPORT_CONFIG_MIN_INT,
-    REPORT_CONFIG_MAX_INT, REPORT_CONFIG_OP, SIGNAL_REMOVE,
-    NO_SENSOR_CLUSTERS, POWER_CONFIGURATION_CHANNEL, BINDABLE_CLUSTERS,
-    DATA_ZHA_GATEWAY, ACCELERATION)
 from .device import ZHADevice, DeviceStatus
-from ..device_entity import ZhaDeviceEntity
 from .channels import (
-    AttributeListeningChannel, EventRelayChannel, ZDOChannel, MAINS_POWERED
+    ZDOChannel, MAINS_POWERED
 )
-from .channels.registry import ZIGBEE_CHANNEL_REGISTRY
 from .helpers import convert_ieee
+from .discovery import (
+    async_process_endpoint, async_dispatch_discovery_info,
+    async_create_device_entity
+)
+from .store import async_get_registry
+from .patches import apply_application_controller_patch
 
 _LOGGER = logging.getLogger(__name__)
 
-SENSOR_TYPES = {}
-BINARY_SENSOR_TYPES = {}
-SMARTTHINGS_HUMIDITY_CLUSTER = 64581
-SMARTTHINGS_ACCELERATION_CLUSTER = 64514
 EntityReference = collections.namedtuple(
     'EntityReference', 'reference_id zha_device cluster_channels device_info')
 
@@ -45,17 +41,52 @@ EntityReference = collections.namedtuple(
 class ZHAGateway:
     """Gateway that handles events that happen on the ZHA Zigbee network."""
 
-    def __init__(self, hass, config, zha_storage):
+    def __init__(self, hass, config):
         """Initialize the gateway."""
         self._hass = hass
         self._config = config
         self._component = EntityComponent(_LOGGER, DOMAIN, hass)
         self._devices = {}
         self._device_registry = collections.defaultdict(list)
-        self.zha_storage = zha_storage
+        self.zha_storage = None
+        self.application_controller = None
+        self.radio_description = None
         hass.data[DATA_ZHA][DATA_ZHA_CORE_COMPONENT] = self._component
         hass.data[DATA_ZHA][DATA_ZHA_GATEWAY] = self
 
+    async def async_initialize(self, config_entry):
+        """Initialize controller and connect radio."""
+        self.zha_storage = await async_get_registry(self._hass)
+
+        usb_path = config_entry.data.get(CONF_USB_PATH)
+        baudrate = self._config.get(CONF_BAUDRATE, DEFAULT_BAUDRATE)
+        radio_type = config_entry.data.get(CONF_RADIO_TYPE)
+
+        radio_details = RADIO_TYPES[radio_type][RADIO]()
+        radio = radio_details[RADIO]
+        self.radio_description = RADIO_TYPES[radio_type][RADIO_DESCRIPTION]
+        await radio.connect(usb_path, baudrate)
+        self._hass.data[DATA_ZHA][DATA_ZHA_RADIO] = radio
+
+        if CONF_DATABASE in self._config:
+            database = self._config[CONF_DATABASE]
+        else:
+            database = os.path.join(
+                self._hass.config.config_dir, DEFAULT_DATABASE_NAME)
+
+        self.application_controller = radio_details[CONTROLLER](
+            radio, database)
+        apply_application_controller_patch(self)
+        self.application_controller.add_listener(self)
+        await self.application_controller.startup(auto_form=True)
+        self._hass.data[DATA_ZHA][DATA_ZHA_BRIDGE_ID] = str(
+            self.application_controller.ieee)
+
+        init_tasks = []
+        for device in self.application_controller.devices.values():
+            init_tasks.append(self.async_device_initialized(device, False))
+        await asyncio.gather(*init_tasks)
+
     def device_joined(self, device):
         """Handle device joined.
 
@@ -166,9 +197,9 @@ class ZHAGateway:
 
         discovery_infos = []
         for endpoint_id, endpoint in device.endpoints.items():
-            self._async_process_endpoint(
+            async_process_endpoint(
-                endpoint_id, endpoint, discovery_infos, device, zha_device,
+                self._hass, self._config, endpoint_id, endpoint,
-                is_new_join
+                discovery_infos, device, zha_device, is_new_join
             )
 
         if is_new_join:
@@ -191,459 +222,11 @@ class ZHAGateway:
             await zha_device.async_initialize(from_cache=True)
 
         for discovery_info in discovery_infos:
-            _async_dispatch_discovery_info(
+            async_dispatch_discovery_info(
                 self._hass,
                 is_new_join,
                 discovery_info
             )
 
-        device_entity = _async_create_device_entity(zha_device)
+        device_entity = async_create_device_entity(zha_device)
         await self._component.async_add_entities([device_entity])
-
-    @callback
-    def _async_process_endpoint(
-            self, endpoint_id, endpoint, discovery_infos, device, zha_device,
-            is_new_join):
-        """Process an endpoint on a zigpy device."""
-        import zigpy.profiles
-
-        if endpoint_id == 0:  # ZDO
-            _async_create_cluster_channel(
-                endpoint,
-                zha_device,
-                is_new_join,
-                channel_class=ZDOChannel
-            )
-            return
-
-        component = None
-        profile_clusters = ([], [])
-        device_key = "{}-{}".format(device.ieee, endpoint_id)
-        node_config = {}
-        if CONF_DEVICE_CONFIG in self._config:
-            node_config = self._config[CONF_DEVICE_CONFIG].get(
-                device_key, {}
-            )
-
-        if endpoint.profile_id in zigpy.profiles.PROFILES:
-            profile = zigpy.profiles.PROFILES[endpoint.profile_id]
-            if zha_const.DEVICE_CLASS.get(endpoint.profile_id,
-                                          {}).get(endpoint.device_type,
-                                                  None):
-                profile_clusters = profile.CLUSTERS[endpoint.device_type]
-                profile_info = zha_const.DEVICE_CLASS[endpoint.profile_id]
-                component = profile_info[endpoint.device_type]
-
-        if ha_const.CONF_TYPE in node_config:
-            component = node_config[ha_const.CONF_TYPE]
-            profile_clusters = zha_const.COMPONENT_CLUSTERS[component]
-
-        if component and component in COMPONENTS:
-            profile_match = _async_handle_profile_match(
-                self._hass, endpoint, profile_clusters, zha_device,
-                component, device_key, is_new_join)
-            discovery_infos.append(profile_match)
-
-        discovery_infos.extend(_async_handle_single_cluster_matches(
-            self._hass,
-            endpoint,
-            zha_device,
-            profile_clusters,
-            device_key,
-            is_new_join
-        ))
-
-
-@callback
-def _async_create_cluster_channel(cluster, zha_device, is_new_join,
-                                  channels=None, channel_class=None):
-    """Create a cluster channel and attach it to a device."""
-    if channel_class is None:
-        channel_class = ZIGBEE_CHANNEL_REGISTRY.get(cluster.cluster_id,
-                                                    AttributeListeningChannel)
-    channel = channel_class(cluster, zha_device)
-    zha_device.add_cluster_channel(channel)
-    if channels is not None:
-        channels.append(channel)
-
-
-@callback
-def _async_dispatch_discovery_info(hass, is_new_join, discovery_info):
-    """Dispatch or store discovery information."""
-    if not discovery_info['channels']:
-        _LOGGER.warning(
-            "there are no channels in the discovery info: %s", discovery_info)
-        return
-    component = discovery_info['component']
-    if is_new_join:
-        async_dispatcher_send(
-            hass,
-            ZHA_DISCOVERY_NEW.format(component),
-            discovery_info
-        )
-    else:
-        hass.data[DATA_ZHA][component][discovery_info['unique_id']] = \
-            discovery_info
-
-
-@callback
-def _async_handle_profile_match(hass, endpoint, profile_clusters, zha_device,
-                                component, device_key, is_new_join):
-    """Dispatch a profile match to the appropriate HA component."""
-    in_clusters = [endpoint.in_clusters[c]
-                   for c in profile_clusters[0]
-                   if c in endpoint.in_clusters]
-    out_clusters = [endpoint.out_clusters[c]
-                    for c in profile_clusters[1]
-                    if c in endpoint.out_clusters]
-
-    channels = []
-
-    for cluster in in_clusters:
-        _async_create_cluster_channel(
-            cluster, zha_device, is_new_join, channels=channels)
-
-    for cluster in out_clusters:
-        _async_create_cluster_channel(
-            cluster, zha_device, is_new_join, channels=channels)
-
-    discovery_info = {
-        'unique_id': device_key,
-        'zha_device': zha_device,
-        'channels': channels,
-        'component': component
-    }
-
-    if component == 'binary_sensor':
-        discovery_info.update({SENSOR_TYPE: UNKNOWN})
-        cluster_ids = []
-        cluster_ids.extend(profile_clusters[0])
-        cluster_ids.extend(profile_clusters[1])
-        for cluster_id in cluster_ids:
-            if cluster_id in BINARY_SENSOR_TYPES:
-                discovery_info.update({
-                    SENSOR_TYPE: BINARY_SENSOR_TYPES.get(
-                        cluster_id, UNKNOWN)
-                })
-                break
-
-    return discovery_info
-
-
-@callback
-def _async_handle_single_cluster_matches(hass, endpoint, zha_device,
-                                         profile_clusters, device_key,
-                                         is_new_join):
-    """Dispatch single cluster matches to HA components."""
-    cluster_matches = []
-    cluster_match_results = []
-    for cluster in endpoint.in_clusters.values():
-        # don't let profiles prevent these channels from being created
-        if cluster.cluster_id in NO_SENSOR_CLUSTERS:
-            cluster_match_results.append(
-                _async_handle_channel_only_cluster_match(
-                    zha_device,
-                    cluster,
-                    is_new_join,
-                ))
-
-        if cluster.cluster_id not in profile_clusters[0]:
-            cluster_match_results.append(_async_handle_single_cluster_match(
-                hass,
-                zha_device,
-                cluster,
-                device_key,
-                zha_const.SINGLE_INPUT_CLUSTER_DEVICE_CLASS,
-                is_new_join,
-            ))
-
-    for cluster in endpoint.out_clusters.values():
-        if cluster.cluster_id not in profile_clusters[1]:
-            cluster_match_results.append(_async_handle_single_cluster_match(
-                hass,
-                zha_device,
-                cluster,
-                device_key,
-                zha_const.SINGLE_OUTPUT_CLUSTER_DEVICE_CLASS,
-                is_new_join,
-            ))
-
-        if cluster.cluster_id in EVENT_RELAY_CLUSTERS:
-            _async_create_cluster_channel(
-                cluster,
-                zha_device,
-                is_new_join,
-                channel_class=EventRelayChannel
-            )
-
-    for cluster_match in cluster_match_results:
-        if cluster_match is not None:
-            cluster_matches.append(cluster_match)
-    return cluster_matches
-
-
-@callback
-def _async_handle_channel_only_cluster_match(
-        zha_device, cluster, is_new_join):
-    """Handle a channel only cluster match."""
-    _async_create_cluster_channel(cluster, zha_device, is_new_join)
-
-
-@callback
-def _async_handle_single_cluster_match(hass, zha_device, cluster, device_key,
-                                       device_classes, is_new_join):
-    """Dispatch a single cluster match to a HA component."""
-    component = None  # sub_component = None
-    for cluster_type, candidate_component in device_classes.items():
-        if isinstance(cluster_type, int):
-            if cluster.cluster_id == cluster_type:
-                component = candidate_component
-        elif isinstance(cluster, cluster_type):
-            component = candidate_component
-            break
-
-    if component is None or component not in COMPONENTS:
-        return
-    channels = []
-    _async_create_cluster_channel(cluster, zha_device, is_new_join,
-                                  channels=channels)
-
-    cluster_key = "{}-{}".format(device_key, cluster.cluster_id)
-    discovery_info = {
-        'unique_id': cluster_key,
-        'zha_device': zha_device,
-        'channels': channels,
-        'entity_suffix': '_{}'.format(cluster.cluster_id),
-        'component': component
-    }
-
-    if component == 'sensor':
-        discovery_info.update({
-            SENSOR_TYPE: SENSOR_TYPES.get(cluster.cluster_id, GENERIC)
-        })
-    if component == 'binary_sensor':
-        discovery_info.update({
-            SENSOR_TYPE: BINARY_SENSOR_TYPES.get(cluster.cluster_id, UNKNOWN)
-        })
-
-    return discovery_info
-
-
-@callback
-def _async_create_device_entity(zha_device):
-    """Create ZHADeviceEntity."""
-    device_entity_channels = []
-    if POWER_CONFIGURATION_CHANNEL in zha_device.cluster_channels:
-        channel = zha_device.cluster_channels.get(POWER_CONFIGURATION_CHANNEL)
-        device_entity_channels.append(channel)
-    return ZhaDeviceEntity(zha_device, device_entity_channels)
-
-
-def establish_device_mappings():
-    """Establish mappings between ZCL objects and HA ZHA objects.
-
-    These cannot be module level, as importing bellows must be done in a
-    in a function.
-    """
-    from zigpy import zcl
-    from zigpy.profiles import PROFILES, zha, zll
-
-    if zha.PROFILE_ID not in DEVICE_CLASS:
-        DEVICE_CLASS[zha.PROFILE_ID] = {}
-    if zll.PROFILE_ID not in DEVICE_CLASS:
-        DEVICE_CLASS[zll.PROFILE_ID] = {}
-
-    EVENT_RELAY_CLUSTERS.append(zcl.clusters.general.LevelControl.cluster_id)
-    EVENT_RELAY_CLUSTERS.append(zcl.clusters.general.OnOff.cluster_id)
-
-    NO_SENSOR_CLUSTERS.append(zcl.clusters.general.Basic.cluster_id)
-    NO_SENSOR_CLUSTERS.append(
-        zcl.clusters.general.PowerConfiguration.cluster_id)
-    NO_SENSOR_CLUSTERS.append(zcl.clusters.lightlink.LightLink.cluster_id)
-
-    BINDABLE_CLUSTERS.append(zcl.clusters.general.LevelControl.cluster_id)
-    BINDABLE_CLUSTERS.append(zcl.clusters.general.OnOff.cluster_id)
-    BINDABLE_CLUSTERS.append(zcl.clusters.lighting.Color.cluster_id)
-
-    DEVICE_CLASS[zha.PROFILE_ID].update({
-        zha.DeviceType.ON_OFF_SWITCH: 'binary_sensor',
-        zha.DeviceType.LEVEL_CONTROL_SWITCH: 'binary_sensor',
-        zha.DeviceType.REMOTE_CONTROL: 'binary_sensor',
-        zha.DeviceType.SMART_PLUG: 'switch',
-        zha.DeviceType.LEVEL_CONTROLLABLE_OUTPUT: 'light',
-        zha.DeviceType.ON_OFF_LIGHT: 'light',
-        zha.DeviceType.DIMMABLE_LIGHT: 'light',
-        zha.DeviceType.COLOR_DIMMABLE_LIGHT: 'light',
-        zha.DeviceType.ON_OFF_LIGHT_SWITCH: 'binary_sensor',
-        zha.DeviceType.DIMMER_SWITCH: 'binary_sensor',
-        zha.DeviceType.COLOR_DIMMER_SWITCH: 'binary_sensor',
-    })
-
-    DEVICE_CLASS[zll.PROFILE_ID].update({
-        zll.DeviceType.ON_OFF_LIGHT: 'light',
-        zll.DeviceType.ON_OFF_PLUGIN_UNIT: 'switch',
-        zll.DeviceType.DIMMABLE_LIGHT: 'light',
-        zll.DeviceType.DIMMABLE_PLUGIN_UNIT: 'light',
-        zll.DeviceType.COLOR_LIGHT: 'light',
-        zll.DeviceType.EXTENDED_COLOR_LIGHT: 'light',
-        zll.DeviceType.COLOR_TEMPERATURE_LIGHT: 'light',
-        zll.DeviceType.COLOR_CONTROLLER: 'binary_sensor',
-        zll.DeviceType.COLOR_SCENE_CONTROLLER: 'binary_sensor',
-        zll.DeviceType.CONTROLLER: 'binary_sensor',
-        zll.DeviceType.SCENE_CONTROLLER: 'binary_sensor',
-        zll.DeviceType.ON_OFF_SENSOR: 'binary_sensor',
-    })
-
-    SINGLE_INPUT_CLUSTER_DEVICE_CLASS.update({
-        zcl.clusters.general.OnOff: 'switch',
-        zcl.clusters.measurement.RelativeHumidity: 'sensor',
-        # this works for now but if we hit conflicts we can break it out to
-        # a different dict that is keyed by manufacturer
-        SMARTTHINGS_HUMIDITY_CLUSTER: 'sensor',
-        zcl.clusters.measurement.TemperatureMeasurement: 'sensor',
-        zcl.clusters.measurement.PressureMeasurement: 'sensor',
-        zcl.clusters.measurement.IlluminanceMeasurement: 'sensor',
-        zcl.clusters.smartenergy.Metering: 'sensor',
-        zcl.clusters.homeautomation.ElectricalMeasurement: 'sensor',
-        zcl.clusters.security.IasZone: 'binary_sensor',
-        zcl.clusters.measurement.OccupancySensing: 'binary_sensor',
-        zcl.clusters.hvac.Fan: 'fan',
-        SMARTTHINGS_ACCELERATION_CLUSTER: 'binary_sensor',
-    })
-
-    SINGLE_OUTPUT_CLUSTER_DEVICE_CLASS.update({
-        zcl.clusters.general.OnOff: 'binary_sensor',
-    })
-
-    SENSOR_TYPES.update({
-        zcl.clusters.measurement.RelativeHumidity.cluster_id: HUMIDITY,
-        SMARTTHINGS_HUMIDITY_CLUSTER: HUMIDITY,
-        zcl.clusters.measurement.TemperatureMeasurement.cluster_id:
-        TEMPERATURE,
-        zcl.clusters.measurement.PressureMeasurement.cluster_id: PRESSURE,
-        zcl.clusters.measurement.IlluminanceMeasurement.cluster_id:
-        ILLUMINANCE,
-        zcl.clusters.smartenergy.Metering.cluster_id: METERING,
-        zcl.clusters.homeautomation.ElectricalMeasurement.cluster_id:
-        ELECTRICAL_MEASUREMENT,
-    })
-
-    BINARY_SENSOR_TYPES.update({
-        zcl.clusters.measurement.OccupancySensing.cluster_id: OCCUPANCY,
-        zcl.clusters.security.IasZone.cluster_id: ZONE,
-        zcl.clusters.general.OnOff.cluster_id: OPENING,
-        SMARTTHINGS_ACCELERATION_CLUSTER: ACCELERATION,
-    })
-
-    CLUSTER_REPORT_CONFIGS.update({
-        zcl.clusters.general.Alarms.cluster_id: [],
-        zcl.clusters.general.Basic.cluster_id: [],
-        zcl.clusters.general.Commissioning.cluster_id: [],
-        zcl.clusters.general.Identify.cluster_id: [],
-        zcl.clusters.general.Groups.cluster_id: [],
-        zcl.clusters.general.Scenes.cluster_id: [],
-        zcl.clusters.general.Partition.cluster_id: [],
-        zcl.clusters.general.Ota.cluster_id: [],
-        zcl.clusters.general.PowerProfile.cluster_id: [],
-        zcl.clusters.general.ApplianceControl.cluster_id: [],
-        zcl.clusters.general.PollControl.cluster_id: [],
-        zcl.clusters.general.GreenPowerProxy.cluster_id: [],
-        zcl.clusters.general.OnOffConfiguration.cluster_id: [],
-        zcl.clusters.lightlink.LightLink.cluster_id: [],
-        zcl.clusters.general.OnOff.cluster_id: [{
-            'attr': 'on_off',
-            'config': REPORT_CONFIG_IMMEDIATE
-        }],
-        zcl.clusters.general.LevelControl.cluster_id: [{
-            'attr': 'current_level',
-            'config': REPORT_CONFIG_ASAP
-        }],
-        zcl.clusters.lighting.Color.cluster_id: [{
-            'attr': 'current_x',
-            'config': REPORT_CONFIG_DEFAULT
-        }, {
-            'attr': 'current_y',
-            'config': REPORT_CONFIG_DEFAULT
-        }, {
-            'attr': 'color_temperature',
-            'config': REPORT_CONFIG_DEFAULT
-        }],
-        zcl.clusters.measurement.RelativeHumidity.cluster_id: [{
-            'attr': 'measured_value',
-            'config': (
-                REPORT_CONFIG_MIN_INT,
-                REPORT_CONFIG_MAX_INT,
-                50
-            )
-        }],
-        zcl.clusters.measurement.TemperatureMeasurement.cluster_id: [{
-            'attr': 'measured_value',
-            'config': (
-                REPORT_CONFIG_MIN_INT,
-                REPORT_CONFIG_MAX_INT,
-                50
-            )
-        }],
-        SMARTTHINGS_ACCELERATION_CLUSTER: [{
-            'attr': 'acceleration',
-            'config': REPORT_CONFIG_ASAP
-        }, {
-            'attr': 'x_axis',
-            'config': REPORT_CONFIG_ASAP
-        }, {
-            'attr': 'y_axis',
-            'config': REPORT_CONFIG_ASAP
-        }, {
-            'attr': 'z_axis',
-            'config': REPORT_CONFIG_ASAP
-        }],
-        SMARTTHINGS_HUMIDITY_CLUSTER: [{
-            'attr': 'measured_value',
-            'config': (
-                REPORT_CONFIG_MIN_INT,
-                REPORT_CONFIG_MAX_INT,
-                50
-            )
-        }],
-        zcl.clusters.measurement.PressureMeasurement.cluster_id: [{
-            'attr': 'measured_value',
-            'config': REPORT_CONFIG_DEFAULT
-        }],
-        zcl.clusters.measurement.IlluminanceMeasurement.cluster_id: [{
-            'attr': 'measured_value',
-            'config': REPORT_CONFIG_DEFAULT
-        }],
-        zcl.clusters.smartenergy.Metering.cluster_id: [{
-            'attr': 'instantaneous_demand',
-            'config': REPORT_CONFIG_DEFAULT
-        }],
-        zcl.clusters.homeautomation.ElectricalMeasurement.cluster_id: [{
-            'attr': 'active_power',
-            'config': REPORT_CONFIG_DEFAULT
-        }],
-        zcl.clusters.general.PowerConfiguration.cluster_id: [{
-            'attr': 'battery_voltage',
-            'config': REPORT_CONFIG_DEFAULT
-        }, {
-            'attr': 'battery_percentage_remaining',
-            'config': REPORT_CONFIG_DEFAULT
-        }],
-        zcl.clusters.measurement.OccupancySensing.cluster_id: [{
-            'attr': 'occupancy',
-            'config': REPORT_CONFIG_IMMEDIATE
-        }],
-        zcl.clusters.hvac.Fan.cluster_id: [{
-            'attr': 'fan_mode',
-            'config': REPORT_CONFIG_OP
-        }],
-    })
-
-    # A map of hass components to all Zigbee clusters it could use
-    for profile_id, classes in DEVICE_CLASS.items():
-        profile = PROFILES[profile_id]
-        for device_type, component in classes.items():
-            if component not in COMPONENT_CLUSTERS:
-                COMPONENT_CLUSTERS[component] = (set(), set())
-            clusters = profile.CLUSTERS[device_type]
-            COMPONENT_CLUSTERS[component][0].update(clusters[0])
-            COMPONENT_CLUSTERS[component][1].update(clusters[1])

homeassistant/components/zha/core/patches.py

@@ -0,0 +1,41 @@
+"""
+Patch functions for Zigbee Home Automation.
+
+For more details about this component, please refer to the documentation at
+https://home-assistant.io/components/zha/
+"""
+
+import types
+
+
+def apply_cluster_listener_patch():
+    """Apply patches to ZHA objects."""
+    # patch zigpy listener to prevent flooding logs with warnings due to
+    # how zigpy implemented its listeners
+    from zigpy.appdb import ClusterPersistingListener
+
+    def zha_send_event(self, cluster, command, args):
+        pass
+
+    ClusterPersistingListener.zha_send_event = types.MethodType(
+        zha_send_event,
+        ClusterPersistingListener
+    )
+
+
+def apply_application_controller_patch(zha_gateway):
+    """Apply patches to ZHA objects."""
+    # Patch handle_message until zigpy can provide an event here
+    def handle_message(sender, is_reply, profile, cluster,
+                       src_ep, dst_ep, tsn, command_id, args):
+        """Handle message from a device."""
+        if not sender.initializing and sender.ieee in zha_gateway.devices and \
+                not zha_gateway.devices[sender.ieee].available:
+            zha_gateway.async_device_became_available(
+                sender, is_reply, profile, cluster, src_ep, dst_ep, tsn,
+                command_id, args
+            )
+        return sender.handle_message(
+            is_reply, profile, cluster, src_ep, dst_ep, tsn, command_id, args)
+
+    zha_gateway.application_controller.handle_message = handle_message

homeassistant/components/zha/core/registries.py

@@ -0,0 +1,271 @@
+"""
+Mapping registries for Zigbee Home Automation.
+
+For more details about this component, please refer to the documentation at
+https://home-assistant.io/components/zha/
+"""
+
+from .const import (
+    DEVICE_CLASS, SINGLE_INPUT_CLUSTER_DEVICE_CLASS,
+    SINGLE_OUTPUT_CLUSTER_DEVICE_CLASS, COMPONENT_CLUSTERS, HUMIDITY,
+    TEMPERATURE, ILLUMINANCE, PRESSURE, METERING, ELECTRICAL_MEASUREMENT,
+    EVENT_RELAY_CLUSTERS, OPENING, ZONE,
+    OCCUPANCY, CLUSTER_REPORT_CONFIGS, REPORT_CONFIG_IMMEDIATE,
+    REPORT_CONFIG_ASAP, REPORT_CONFIG_DEFAULT, REPORT_CONFIG_MIN_INT,
+    REPORT_CONFIG_MAX_INT, REPORT_CONFIG_OP,
+    NO_SENSOR_CLUSTERS, BINDABLE_CLUSTERS, ACCELERATION, SENSOR_TYPES,
+    BINARY_SENSOR_TYPES, RADIO_TYPES, RadioType, RADIO, RADIO_DESCRIPTION,
+    CONTROLLER
+)
+
+SMARTTHINGS_HUMIDITY_CLUSTER = 64581
+SMARTTHINGS_ACCELERATION_CLUSTER = 64514
+
+
+def establish_device_mappings():
+    """Establish mappings between ZCL objects and HA ZHA objects.
+
+    These cannot be module level, as importing bellows must be done in a
+    in a function.
+    """
+    from zigpy import zcl
+    from zigpy.profiles import PROFILES, zha, zll
+
+    if zha.PROFILE_ID not in DEVICE_CLASS:
+        DEVICE_CLASS[zha.PROFILE_ID] = {}
+    if zll.PROFILE_ID not in DEVICE_CLASS:
+        DEVICE_CLASS[zll.PROFILE_ID] = {}
+
+    def get_ezsp_radio():
+        import bellows.ezsp
+        from bellows.zigbee.application import ControllerApplication
+        return {
+            RADIO: bellows.ezsp.EZSP(),
+            CONTROLLER: ControllerApplication
+        }
+
+    RADIO_TYPES[RadioType.ezsp.name] = {
+        RADIO: get_ezsp_radio,
+        RADIO_DESCRIPTION: 'EZSP'
+    }
+
+    def get_xbee_radio():
+        import zigpy_xbee.api
+        from zigpy_xbee.zigbee.application import ControllerApplication
+        return {
+            RADIO: zigpy_xbee.api.XBee(),
+            CONTROLLER: ControllerApplication
+        }
+
+    RADIO_TYPES[RadioType.xbee.name] = {
+        RADIO: get_xbee_radio,
+        RADIO_DESCRIPTION: 'XBee'
+    }
+
+    def get_deconz_radio():
+        import zigpy_deconz.api
+        from zigpy_deconz.zigbee.application import ControllerApplication
+        return {
+            RADIO: zigpy_deconz.api.Deconz(),
+            CONTROLLER: ControllerApplication
+        }
+
+    RADIO_TYPES[RadioType.deconz.name] = {
+        RADIO: get_deconz_radio,
+        RADIO_DESCRIPTION: 'Deconz'
+    }
+
+    EVENT_RELAY_CLUSTERS.append(zcl.clusters.general.LevelControl.cluster_id)
+    EVENT_RELAY_CLUSTERS.append(zcl.clusters.general.OnOff.cluster_id)
+
+    NO_SENSOR_CLUSTERS.append(zcl.clusters.general.Basic.cluster_id)
+    NO_SENSOR_CLUSTERS.append(
+        zcl.clusters.general.PowerConfiguration.cluster_id)
+    NO_SENSOR_CLUSTERS.append(zcl.clusters.lightlink.LightLink.cluster_id)
+
+    BINDABLE_CLUSTERS.append(zcl.clusters.general.LevelControl.cluster_id)
+    BINDABLE_CLUSTERS.append(zcl.clusters.general.OnOff.cluster_id)
+    BINDABLE_CLUSTERS.append(zcl.clusters.lighting.Color.cluster_id)
+
+    DEVICE_CLASS[zha.PROFILE_ID].update({
+        zha.DeviceType.ON_OFF_SWITCH: 'binary_sensor',
+        zha.DeviceType.LEVEL_CONTROL_SWITCH: 'binary_sensor',
+        zha.DeviceType.REMOTE_CONTROL: 'binary_sensor',
+        zha.DeviceType.SMART_PLUG: 'switch',
+        zha.DeviceType.LEVEL_CONTROLLABLE_OUTPUT: 'light',
+        zha.DeviceType.ON_OFF_LIGHT: 'light',
+        zha.DeviceType.DIMMABLE_LIGHT: 'light',
+        zha.DeviceType.COLOR_DIMMABLE_LIGHT: 'light',
+        zha.DeviceType.ON_OFF_LIGHT_SWITCH: 'binary_sensor',
+        zha.DeviceType.DIMMER_SWITCH: 'binary_sensor',
+        zha.DeviceType.COLOR_DIMMER_SWITCH: 'binary_sensor',
+    })
+
+    DEVICE_CLASS[zll.PROFILE_ID].update({
+        zll.DeviceType.ON_OFF_LIGHT: 'light',
+        zll.DeviceType.ON_OFF_PLUGIN_UNIT: 'switch',
+        zll.DeviceType.DIMMABLE_LIGHT: 'light',
+        zll.DeviceType.DIMMABLE_PLUGIN_UNIT: 'light',
+        zll.DeviceType.COLOR_LIGHT: 'light',
+        zll.DeviceType.EXTENDED_COLOR_LIGHT: 'light',
+        zll.DeviceType.COLOR_TEMPERATURE_LIGHT: 'light',
+        zll.DeviceType.COLOR_CONTROLLER: 'binary_sensor',
+        zll.DeviceType.COLOR_SCENE_CONTROLLER: 'binary_sensor',
+        zll.DeviceType.CONTROLLER: 'binary_sensor',
+        zll.DeviceType.SCENE_CONTROLLER: 'binary_sensor',
+        zll.DeviceType.ON_OFF_SENSOR: 'binary_sensor',
+    })
+
+    SINGLE_INPUT_CLUSTER_DEVICE_CLASS.update({
+        zcl.clusters.general.OnOff: 'switch',
+        zcl.clusters.measurement.RelativeHumidity: 'sensor',
+        # this works for now but if we hit conflicts we can break it out to
+        # a different dict that is keyed by manufacturer
+        SMARTTHINGS_HUMIDITY_CLUSTER: 'sensor',
+        zcl.clusters.measurement.TemperatureMeasurement: 'sensor',
+        zcl.clusters.measurement.PressureMeasurement: 'sensor',
+        zcl.clusters.measurement.IlluminanceMeasurement: 'sensor',
+        zcl.clusters.smartenergy.Metering: 'sensor',
+        zcl.clusters.homeautomation.ElectricalMeasurement: 'sensor',
+        zcl.clusters.security.IasZone: 'binary_sensor',
+        zcl.clusters.measurement.OccupancySensing: 'binary_sensor',
+        zcl.clusters.hvac.Fan: 'fan',
+        SMARTTHINGS_ACCELERATION_CLUSTER: 'binary_sensor',
+    })
+
+    SINGLE_OUTPUT_CLUSTER_DEVICE_CLASS.update({
+        zcl.clusters.general.OnOff: 'binary_sensor',
+    })
+
+    SENSOR_TYPES.update({
+        zcl.clusters.measurement.RelativeHumidity.cluster_id: HUMIDITY,
+        SMARTTHINGS_HUMIDITY_CLUSTER: HUMIDITY,
+        zcl.clusters.measurement.TemperatureMeasurement.cluster_id:
+        TEMPERATURE,
+        zcl.clusters.measurement.PressureMeasurement.cluster_id: PRESSURE,
+        zcl.clusters.measurement.IlluminanceMeasurement.cluster_id:
+        ILLUMINANCE,
+        zcl.clusters.smartenergy.Metering.cluster_id: METERING,
+        zcl.clusters.homeautomation.ElectricalMeasurement.cluster_id:
+        ELECTRICAL_MEASUREMENT,
+    })
+
+    BINARY_SENSOR_TYPES.update({
+        zcl.clusters.measurement.OccupancySensing.cluster_id: OCCUPANCY,
+        zcl.clusters.security.IasZone.cluster_id: ZONE,
+        zcl.clusters.general.OnOff.cluster_id: OPENING,
+        SMARTTHINGS_ACCELERATION_CLUSTER: ACCELERATION,
+    })
+
+    CLUSTER_REPORT_CONFIGS.update({
+        zcl.clusters.general.Alarms.cluster_id: [],
+        zcl.clusters.general.Basic.cluster_id: [],
+        zcl.clusters.general.Commissioning.cluster_id: [],
+        zcl.clusters.general.Identify.cluster_id: [],
+        zcl.clusters.general.Groups.cluster_id: [],
+        zcl.clusters.general.Scenes.cluster_id: [],
+        zcl.clusters.general.Partition.cluster_id: [],
+        zcl.clusters.general.Ota.cluster_id: [],
+        zcl.clusters.general.PowerProfile.cluster_id: [],
+        zcl.clusters.general.ApplianceControl.cluster_id: [],
+        zcl.clusters.general.PollControl.cluster_id: [],
+        zcl.clusters.general.GreenPowerProxy.cluster_id: [],
+        zcl.clusters.general.OnOffConfiguration.cluster_id: [],
+        zcl.clusters.lightlink.LightLink.cluster_id: [],
+        zcl.clusters.general.OnOff.cluster_id: [{
+            'attr': 'on_off',
+            'config': REPORT_CONFIG_IMMEDIATE
+        }],
+        zcl.clusters.general.LevelControl.cluster_id: [{
+            'attr': 'current_level',
+            'config': REPORT_CONFIG_ASAP
+        }],
+        zcl.clusters.lighting.Color.cluster_id: [{
+            'attr': 'current_x',
+            'config': REPORT_CONFIG_DEFAULT
+        }, {
+            'attr': 'current_y',
+            'config': REPORT_CONFIG_DEFAULT
+        }, {
+            'attr': 'color_temperature',
+            'config': REPORT_CONFIG_DEFAULT
+        }],
+        zcl.clusters.measurement.RelativeHumidity.cluster_id: [{
+            'attr': 'measured_value',
+            'config': (
+                REPORT_CONFIG_MIN_INT,
+                REPORT_CONFIG_MAX_INT,
+                50
+            )
+        }],
+        zcl.clusters.measurement.TemperatureMeasurement.cluster_id: [{
+            'attr': 'measured_value',
+            'config': (
+                REPORT_CONFIG_MIN_INT,
+                REPORT_CONFIG_MAX_INT,
+                50
+            )
+        }],
+        SMARTTHINGS_ACCELERATION_CLUSTER: [{
+            'attr': 'acceleration',
+            'config': REPORT_CONFIG_ASAP
+        }, {
+            'attr': 'x_axis',
+            'config': REPORT_CONFIG_ASAP
+        }, {
+            'attr': 'y_axis',
+            'config': REPORT_CONFIG_ASAP
+        }, {
+            'attr': 'z_axis',
+            'config': REPORT_CONFIG_ASAP
+        }],
+        SMARTTHINGS_HUMIDITY_CLUSTER: [{
+            'attr': 'measured_value',
+            'config': (
+                REPORT_CONFIG_MIN_INT,
+                REPORT_CONFIG_MAX_INT,
+                50
+            )
+        }],
+        zcl.clusters.measurement.PressureMeasurement.cluster_id: [{
+            'attr': 'measured_value',
+            'config': REPORT_CONFIG_DEFAULT
+        }],
+        zcl.clusters.measurement.IlluminanceMeasurement.cluster_id: [{
+            'attr': 'measured_value',
+            'config': REPORT_CONFIG_DEFAULT
+        }],
+        zcl.clusters.smartenergy.Metering.cluster_id: [{
+            'attr': 'instantaneous_demand',
+            'config': REPORT_CONFIG_DEFAULT
+        }],
+        zcl.clusters.homeautomation.ElectricalMeasurement.cluster_id: [{
+            'attr': 'active_power',
+            'config': REPORT_CONFIG_DEFAULT
+        }],
+        zcl.clusters.general.PowerConfiguration.cluster_id: [{
+            'attr': 'battery_voltage',
+            'config': REPORT_CONFIG_DEFAULT
+        }, {
+            'attr': 'battery_percentage_remaining',
+            'config': REPORT_CONFIG_DEFAULT
+        }],
+        zcl.clusters.measurement.OccupancySensing.cluster_id: [{
+            'attr': 'occupancy',
+            'config': REPORT_CONFIG_IMMEDIATE
+        }],
+        zcl.clusters.hvac.Fan.cluster_id: [{
+            'attr': 'fan_mode',
+            'config': REPORT_CONFIG_OP
+        }],
+    })
+
+    # A map of hass components to all Zigbee clusters it could use
+    for profile_id, classes in DEVICE_CLASS.items():
+        profile = PROFILES[profile_id]
+        for device_type, component in classes.items():
+            if component not in COMPONENT_CLUSTERS:
+                COMPONENT_CLUSTERS[component] = (set(), set())
+            clusters = profile.CLUSTERS[device_type]
+            COMPONENT_CLUSTERS[component][0].update(clusters[0])
+            COMPONENT_CLUSTERS[component][1].update(clusters[1])

tests/components/zha/conftest.py

@@ -6,7 +6,8 @@ from homeassistant.components.zha.core.const import (
     DOMAIN, DATA_ZHA, COMPONENTS
 )
 from homeassistant.components.zha.core.gateway import ZHAGateway
-from homeassistant.components.zha.core.gateway import establish_device_mappings
+from homeassistant.components.zha.core.registries import \
+    establish_device_mappings
 from homeassistant.components.zha.core.channels.registry \
     import populate_channel_registry
 from .common import async_setup_entry
@@ -36,7 +37,9 @@ async def zha_gateway_fixture(hass):
             hass.data[DATA_ZHA].get(component, {})
         )
     zha_storage = await async_get_registry(hass)
-    return ZHAGateway(hass, {}, zha_storage)
+    gateway = ZHAGateway(hass, {})
+    gateway.zha_storage = zha_storage
+    return gateway
 
 
 @pytest.fixture(autouse=True)

tests/components/zha/test_api.py

@@ -1,5 +1,4 @@
 """Test ZHA API."""
-from unittest.mock import Mock
 import pytest
 from homeassistant.components.switch import DOMAIN
 from homeassistant.components.zha.api import (
@@ -18,7 +17,7 @@ async def zha_client(hass, config_entry, zha_gateway, hass_ws_client):
     from zigpy.zcl.clusters.general import OnOff, Basic
 
     # load the ZHA API
-    async_load_api(hass, Mock(), zha_gateway)
+    async_load_api(hass)
 
     # create zigpy device
     await async_init_zigpy_device(