Refactor ZHA config flow (#35397)

* Refactor ZHA config flow

Present more descriptive list of radio types when user has to pick one.

* Update docstring.

* Add some common models to EZSP radio.

* Add more model names

More model names, less english since radio types won't be translated.

homeassistant/components/zha/config_flow.py

@@ -12,11 +12,9 @@ from .core.const import (  # pylint:disable=unused-import
     CONF_BAUDRATE,
     CONF_FLOWCONTROL,
     CONF_RADIO_TYPE,
-    CONTROLLER,
     DOMAIN,
     RadioType,
 )
-from .core.registries import RADIO_TYPES
 
 CONF_MANUAL_PATH = "Enter Manually"
 SUPPORTED_PORT_SETTINGS = (
@@ -75,7 +73,7 @@ class ZhaFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
         """Select radio type."""
 
         if user_input is not None:
-            self._radio_type = user_input[CONF_RADIO_TYPE]
+            self._radio_type = RadioType.get_by_description(user_input[CONF_RADIO_TYPE])
             return await self.async_step_port_config()
 
         schema = {vol.Required(CONF_RADIO_TYPE): vol.In(sorted(RadioType.list()))}
@@ -86,7 +84,7 @@ class ZhaFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
     async def async_step_port_config(self, user_input=None):
         """Enter port settings specific for this type of radio."""
         errors = {}
-        app_cls = RADIO_TYPES[self._radio_type][CONTROLLER]
+        app_cls = RadioType[self._radio_type].controller
 
         if user_input is not None:
             self._device_path = user_input.get(CONF_DEVICE_PATH)
@@ -121,11 +119,10 @@ class ZhaFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
 
 async def detect_radios(dev_path: str) -> Optional[Dict[str, Any]]:
     """Probe all radio types on the device port."""
-    for radio in RadioType.list():
+    for radio in RadioType:
-        app_cls = RADIO_TYPES[radio][CONTROLLER]
+        dev_config = radio.controller.SCHEMA_DEVICE({CONF_DEVICE_PATH: dev_path})
-        dev_config = app_cls.SCHEMA_DEVICE({CONF_DEVICE_PATH: dev_path})
+        if await radio.controller.probe(dev_config):
-        if await app_cls.probe(dev_config):
+            return {CONF_RADIO_TYPE: radio.name, CONF_DEVICE: dev_config}
-            return {CONF_RADIO_TYPE: radio, CONF_DEVICE: dev_config}
 
     return None
 

homeassistant/components/zha/core/const.py

@@ -1,8 +1,14 @@
 """All constants related to the ZHA component."""
 import enum
 import logging
+from typing import List
 
+import bellows.zigbee.application
 from zigpy.config import CONF_DEVICE_PATH  # noqa: F401 # pylint: disable=unused-import
+import zigpy_cc.zigbee.application
+import zigpy_deconz.zigbee.application
+import zigpy_xbee.zigbee.application
+import zigpy_zigate.zigbee.application
 
 from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR
 from homeassistant.components.cover import DOMAIN as COVER
@@ -13,6 +19,8 @@ from homeassistant.components.lock import DOMAIN as LOCK
 from homeassistant.components.sensor import DOMAIN as SENSOR
 from homeassistant.components.switch import DOMAIN as SWITCH
 
+from .typing import CALLABLE_T
+
 ATTR_ARGS = "args"
 ATTR_ATTRIBUTE = "attribute"
 ATTR_ATTRIBUTE_ID = "attribute_id"
@@ -98,7 +106,6 @@ CONF_FLOWCONTROL = "flow_control"
 CONF_RADIO_TYPE = "radio_type"
 CONF_USB_PATH = "usb_path"
 CONF_ZIGPY = "zigpy_config"
-CONTROLLER = "controller"
 
 DATA_DEVICE_CONFIG = "zha_device_config"
 DATA_ZHA = "zha"
@@ -149,16 +156,51 @@ POWER_BATTERY_OR_UNKNOWN = "Battery or Unknown"
 class RadioType(enum.Enum):
     """Possible options for radio type."""
 
-    ezsp = "ezsp"
+    ezsp = (
-    deconz = "deconz"
+        "ESZP: HUSBZB-1, Elelabs, Telegesis, Silabs EmberZNet protocol",
-    ti_cc = "ti_cc"
+        bellows.zigbee.application.ControllerApplication,
-    zigate = "zigate"
+    )
-    xbee = "xbee"
+    deconz = (
+        "Conbee, Conbee II, RaspBee radios from dresden elektronik",
+        zigpy_deconz.zigbee.application.ControllerApplication,
+    )
+    ti_cc = (
+        "TI_CC: CC2531, CC2530, CC2652R, CC1352 etc, Texas Instruments ZNP protocol",
+        zigpy_cc.zigbee.application.ControllerApplication,
+    )
+    zigate = "ZiGate Radio", zigpy_zigate.zigbee.application.ControllerApplication
+    xbee = (
+        "Digi XBee S2C, XBee 3 radios",
+        zigpy_xbee.zigbee.application.ControllerApplication,
+    )
 
     @classmethod
-    def list(cls):
+    def list(cls) -> List[str]:
-        """Return list of enum's values."""
+        """Return a list of descriptions."""
-        return [e.value for e in RadioType]
+        return [e.description for e in RadioType]
+
+    @classmethod
+    def get_by_description(cls, description: str) -> str:
+        """Get radio by description."""
+        for radio in cls:
+            if radio.description == description:
+                return radio.name
+        raise ValueError
+
+    def __init__(self, description: str, controller_cls: CALLABLE_T):
+        """Init instance."""
+        self._desc = description
+        self._ctrl_cls = controller_cls
+
+    @property
+    def controller(self) -> CALLABLE_T:
+        """Return controller class."""
+        return self._ctrl_cls
+
+    @property
+    def description(self) -> str:
+        """Return radio type description."""
+        return self._desc
 
 
 REPORT_CONFIG_MAX_INT = 900
@@ -262,7 +304,6 @@ ZHA_GW_MSG_GROUP_REMOVED = "group_removed"
 ZHA_GW_MSG_LOG_ENTRY = "log_entry"
 ZHA_GW_MSG_LOG_OUTPUT = "log_output"
 ZHA_GW_MSG_RAW_INIT = "raw_device_initialized"
-ZHA_GW_RADIO_DESCRIPTION = "radio_description"
 
 EFFECT_BLINK = 0x00
 EFFECT_BREATHE = 0x01

homeassistant/components/zha/core/gateway.py

@@ -37,7 +37,6 @@ from .const import (
     CONF_DATABASE,
     CONF_RADIO_TYPE,
     CONF_ZIGPY,
-    CONTROLLER,
     DATA_ZHA,
     DATA_ZHA_BRIDGE_ID,
     DATA_ZHA_GATEWAY,
@@ -73,12 +72,12 @@ from .const import (
     ZHA_GW_MSG_LOG_ENTRY,
     ZHA_GW_MSG_LOG_OUTPUT,
     ZHA_GW_MSG_RAW_INIT,
-    ZHA_GW_RADIO_DESCRIPTION,
+    RadioType,
 )
 from .device import DeviceStatus, ZHADevice
 from .group import GroupMember, ZHAGroup
 from .patches import apply_application_controller_patch
-from .registries import GROUP_ENTITY_DOMAINS, RADIO_TYPES
+from .registries import GROUP_ENTITY_DOMAINS
 from .store import async_get_registry
 from .typing import ZhaGroupType, ZigpyEndpointType, ZigpyGroupType
 
@@ -125,8 +124,8 @@ class ZHAGateway:
 
         radio_type = self._config_entry.data[CONF_RADIO_TYPE]
 
-        app_controller_cls = RADIO_TYPES[radio_type][CONTROLLER]
+        app_controller_cls = RadioType[radio_type].controller
-        self.radio_description = RADIO_TYPES[radio_type][ZHA_GW_RADIO_DESCRIPTION]
+        self.radio_description = RadioType[radio_type].description
 
         app_config = self._config.get(CONF_ZIGPY, {})
         database = self._config.get(

homeassistant/components/zha/core/registries.py

@@ -3,14 +3,9 @@ import collections
 from typing import Callable, Dict, List, Set, Tuple, Union
 
 import attr
-import bellows.zigbee.application
 import zigpy.profiles.zha
 import zigpy.profiles.zll
 import zigpy.zcl as zcl
-import zigpy_cc.zigbee.application
-import zigpy_deconz.zigbee.application
-import zigpy_xbee.zigbee.application
-import zigpy_zigate.zigbee.application
 
 from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR
 from homeassistant.components.cover import DOMAIN as COVER
@@ -23,7 +18,6 @@ from homeassistant.components.switch import DOMAIN as SWITCH
 
 # importing channels updates registries
 from . import channels as zha_channels  # noqa: F401 pylint: disable=unused-import
-from .const import CONTROLLER, ZHA_GW_RADIO_DESCRIPTION, RadioType
 from .decorators import CALLABLE_T, DictRegistry, SetRegistry
 from .typing import ChannelType
 
@@ -124,29 +118,6 @@ LIGHT_CLUSTERS = SetRegistry()
 OUTPUT_CHANNEL_ONLY_CLUSTERS = SetRegistry()
 CLIENT_CHANNELS_REGISTRY = DictRegistry()
 
-RADIO_TYPES = {
-    RadioType.deconz.name: {
-        CONTROLLER: zigpy_deconz.zigbee.application.ControllerApplication,
-        ZHA_GW_RADIO_DESCRIPTION: "Deconz",
-    },
-    RadioType.ezsp.name: {
-        CONTROLLER: bellows.zigbee.application.ControllerApplication,
-        ZHA_GW_RADIO_DESCRIPTION: "EZSP",
-    },
-    RadioType.ti_cc.name: {
-        CONTROLLER: zigpy_cc.zigbee.application.ControllerApplication,
-        ZHA_GW_RADIO_DESCRIPTION: "TI CC",
-    },
-    RadioType.xbee.name: {
-        CONTROLLER: zigpy_xbee.zigbee.application.ControllerApplication,
-        ZHA_GW_RADIO_DESCRIPTION: "XBee",
-    },
-    RadioType.zigate.name: {
-        CONTROLLER: zigpy_zigate.zigbee.application.ControllerApplication,
-        ZHA_GW_RADIO_DESCRIPTION: "ZiGate",
-    },
-}
-
 COMPONENT_CLUSTERS = {
     BINARY_SENSOR: BINARY_SENSOR_CLUSTERS,
     DEVICE_TRACKER: DEVICE_TRACKER_CLUSTERS,

tests/components/zha/conftest.py

@@ -1,5 +1,4 @@
 """Test configuration for the ZHA component."""
-from unittest import mock
 
 import pytest
 import zigpy
@@ -10,12 +9,11 @@ import zigpy.types
 
 import homeassistant.components.zha.core.const as zha_const
 import homeassistant.components.zha.core.device as zha_core_device
-import homeassistant.components.zha.core.registries as zha_regs
 from homeassistant.setup import async_setup_component
 
 from .common import FakeDevice, FakeEndpoint, get_zha_gateway
 
-import tests.async_mock
+from tests.async_mock import AsyncMock, MagicMock, PropertyMock, patch
 from tests.common import MockConfigEntry
 
 FIXTURE_GRP_ID = 0x1001
@@ -25,27 +23,19 @@ FIXTURE_GRP_NAME = "fixture group"
 @pytest.fixture
 def zigpy_app_controller():
     """Zigpy ApplicationController fixture."""
-    app = mock.MagicMock(spec_set=ControllerApplication)
+    app = MagicMock(spec_set=ControllerApplication)
-    app.startup = tests.async_mock.AsyncMock()
+    app.startup = AsyncMock()
-    app.shutdown = tests.async_mock.AsyncMock()
+    app.shutdown = AsyncMock()
     groups = zigpy.group.Groups(app)
     groups.add_group(FIXTURE_GRP_ID, FIXTURE_GRP_NAME, suppress_event=True)
     app.configure_mock(groups=groups)
-    type(app).ieee = mock.PropertyMock()
+    type(app).ieee = PropertyMock()
     app.ieee.return_value = zigpy.types.EUI64.convert("00:15:8d:00:02:32:4f:32")
-    type(app).nwk = mock.PropertyMock(return_value=zigpy.types.NWK(0x0000))
+    type(app).nwk = PropertyMock(return_value=zigpy.types.NWK(0x0000))
-    type(app).devices = mock.PropertyMock(return_value={})
+    type(app).devices = PropertyMock(return_value={})
     return app
 
 
-@pytest.fixture
-def zigpy_radio():
-    """Zigpy radio mock."""
-    radio = mock.MagicMock()
-    radio.connect = tests.async_mock.AsyncMock()
-    return radio
-
-
 @pytest.fixture(name="config_entry")
 async def config_entry_fixture(hass):
     """Fixture representing a config entry."""
@@ -54,7 +44,7 @@ async def config_entry_fixture(hass):
         domain=zha_const.DOMAIN,
         data={
             zigpy.config.CONF_DEVICE: {zigpy.config.CONF_DEVICE_PATH: "/dev/ttyUSB0"},
-            zha_const.CONF_RADIO_TYPE: "MockRadio",
+            zha_const.CONF_RADIO_TYPE: "ezsp",
         },
     )
     entry.add_to_hass(hass)
@@ -62,22 +52,18 @@ async def config_entry_fixture(hass):
 
 
 @pytest.fixture
-def setup_zha(hass, config_entry, zigpy_app_controller, zigpy_radio):
+def setup_zha(hass, config_entry, zigpy_app_controller):
     """Set up ZHA component."""
     zha_config = {zha_const.CONF_ENABLE_QUIRKS: False}
-    app_ctrl = mock.MagicMock()
-    app_ctrl.new = tests.async_mock.AsyncMock(return_value=zigpy_app_controller)
-    app_ctrl.SCHEMA = zigpy.config.CONFIG_SCHEMA
-    app_ctrl.SCHEMA_DEVICE = zigpy.config.SCHEMA_DEVICE
 
-    radio_details = {
+    p1 = patch(
-        zha_const.CONTROLLER: app_ctrl,
+        "bellows.zigbee.application.ControllerApplication.new",
-        zha_const.ZHA_GW_RADIO_DESCRIPTION: "mock radio",
+        return_value=zigpy_app_controller,
-    }
+    )
 
     async def _setup(config=None):
         config = config or {}
-        with mock.patch.dict(zha_regs.RADIO_TYPES, {"MockRadio": radio_details}):
+        with p1:
             status = await async_setup_component(
                 hass, zha_const.DOMAIN, {zha_const.DOMAIN: {**zha_config, **config}}
             )
@@ -92,12 +78,12 @@ def channel():
     """Channel mock factory fixture."""
 
     def channel(name: str, cluster_id: int, endpoint_id: int = 1):
-        ch = mock.MagicMock()
+        ch = MagicMock()
         ch.name = name
         ch.generic_id = f"channel_0x{cluster_id:04x}"
         ch.id = f"{endpoint_id}:0x{cluster_id:04x}"
-        ch.async_configure = tests.async_mock.AsyncMock()
+        ch.async_configure = AsyncMock()
-        ch.async_initialize = tests.async_mock.AsyncMock()
+        ch.async_initialize = AsyncMock()
         return ch
 
     return channel
@@ -201,7 +187,7 @@ def zha_device_mock(hass, zigpy_device_mock):
         zigpy_device = zigpy_device_mock(
             endpoints, ieee, manufacturer, model, node_desc
         )
-        zha_device = zha_core_device.ZHADevice(hass, zigpy_device, mock.MagicMock())
+        zha_device = zha_core_device.ZHADevice(hass, zigpy_device, MagicMock())
         return zha_device
 
     return _zha_device

tests/components/zha/test_config_flow.py

@@ -8,8 +8,7 @@ import zigpy.config
 
 from homeassistant import setup
 from homeassistant.components.zha import config_flow
-from homeassistant.components.zha.core.const import CONF_RADIO_TYPE, CONTROLLER, DOMAIN
+from homeassistant.components.zha.core.const import CONF_RADIO_TYPE, DOMAIN, RadioType
-from homeassistant.components.zha.core.registries import RADIO_TYPES
 from homeassistant.config_entries import SOURCE_USER
 from homeassistant.const import CONF_SOURCE
 from homeassistant.data_entry_flow import RESULT_TYPE_CREATE_ENTRY, RESULT_TYPE_FORM
@@ -96,11 +95,12 @@ async def test_user_flow_manual(hass):
     assert result["step_id"] == "pick_radio"
 
 
-async def test_pick_radio_flow(hass):
+@pytest.mark.parametrize("radio_type", RadioType.list())
+async def test_pick_radio_flow(hass, radio_type):
     """Test radio picker."""
 
     result = await hass.config_entries.flow.async_init(
-        DOMAIN, context={CONF_SOURCE: "pick_radio"}, data={CONF_RADIO_TYPE: "ezsp"}
+        DOMAIN, context={CONF_SOURCE: "pick_radio"}, data={CONF_RADIO_TYPE: radio_type}
     )
     assert result["type"] == RESULT_TYPE_FORM
     assert result["step_id"] == "port_config"
@@ -117,15 +117,27 @@ async def test_user_flow_existing_config_entry(hass):
     assert result["type"] == "abort"
 
 
-async def test_probe_radios(hass):
+@patch("zigpy_cc.zigbee.application.ControllerApplication.probe", return_value=False)
+@patch(
+    "zigpy_deconz.zigbee.application.ControllerApplication.probe", return_value=False
+)
+@patch(
+    "zigpy_zigate.zigbee.application.ControllerApplication.probe", return_value=False
+)
+@patch("zigpy_xbee.zigbee.application.ControllerApplication.probe", return_value=False)
+async def test_probe_radios(xbee_probe, zigate_probe, deconz_probe, cc_probe, hass):
     """Test detect radios."""
     app_ctrl_cls = MagicMock()
     app_ctrl_cls.SCHEMA_DEVICE = zigpy.config.SCHEMA_DEVICE
     app_ctrl_cls.probe = AsyncMock(side_effect=(True, False))
 
-    with patch.dict(config_flow.RADIO_TYPES, {"ezsp": {CONTROLLER: app_ctrl_cls}}):
+    p1 = patch(
+        "bellows.zigbee.application.ControllerApplication.probe",
+        side_effect=(True, False),
+    )
+    with p1 as probe_mock:
         res = await config_flow.detect_radios("/dev/null")
-        assert app_ctrl_cls.probe.await_count == 1
+        assert probe_mock.await_count == 1
         assert res[CONF_RADIO_TYPE] == "ezsp"
         assert zigpy.config.CONF_DEVICE in res
         assert (
@@ -134,51 +146,45 @@ async def test_probe_radios(hass):
 
         res = await config_flow.detect_radios("/dev/null")
         assert res is None
+        assert xbee_probe.await_count == 1
+        assert zigate_probe.await_count == 1
+        assert deconz_probe.await_count == 1
+        assert cc_probe.await_count == 1
 
 
-async def test_user_port_config_fail(hass):
+@patch("bellows.zigbee.application.ControllerApplication.probe", return_value=False)
+async def test_user_port_config_fail(probe_mock, hass):
     """Test port config flow."""
-    app_ctrl_cls = MagicMock()
-    app_ctrl_cls.SCHEMA_DEVICE = zigpy.config.SCHEMA_DEVICE
-    app_ctrl_cls.probe = AsyncMock(return_value=False)
 
     result = await hass.config_entries.flow.async_init(
-        DOMAIN, context={CONF_SOURCE: "pick_radio"}, data={CONF_RADIO_TYPE: "ezsp"}
+        DOMAIN,
+        context={CONF_SOURCE: "pick_radio"},
+        data={CONF_RADIO_TYPE: RadioType.ezsp.description},
     )
 
-    with patch.dict(config_flow.RADIO_TYPES, {"ezsp": {CONTROLLER: app_ctrl_cls}}):
     result = await hass.config_entries.flow.async_configure(
-            result["flow_id"],
+        result["flow_id"], user_input={zigpy.config.CONF_DEVICE_PATH: "/dev/ttyUSB33"},
-            user_input={zigpy.config.CONF_DEVICE_PATH: "/dev/ttyUSB33"},
     )
     assert result["type"] == RESULT_TYPE_FORM
     assert result["step_id"] == "port_config"
     assert result["errors"]["base"] == "cannot_connect"
+    assert probe_mock.await_count == 1
 
 
-@pytest.mark.parametrize(
-    "radio_type, orig_ctrl_cls",
-    ((name, r[CONTROLLER]) for name, r in RADIO_TYPES.items()),
-)
 @patch("homeassistant.components.zha.async_setup_entry", AsyncMock(return_value=True))
-async def test_user_port_config(hass, radio_type, orig_ctrl_cls):
+@patch("bellows.zigbee.application.ControllerApplication.probe", return_value=True)
+async def test_user_port_config(probe_mock, hass):
     """Test port config."""
-    app_ctrl_cls = MagicMock()
-    app_ctrl_cls.SCHEMA_DEVICE = orig_ctrl_cls.SCHEMA_DEVICE
-    app_ctrl_cls.probe = AsyncMock(return_value=True)
     await setup.async_setup_component(hass, "persistent_notification", {})
 
     result = await hass.config_entries.flow.async_init(
-        DOMAIN, context={CONF_SOURCE: "pick_radio"}, data={CONF_RADIO_TYPE: radio_type}
+        DOMAIN,
+        context={CONF_SOURCE: "pick_radio"},
+        data={CONF_RADIO_TYPE: RadioType.ezsp.description},
     )
 
-    with patch.dict(
-        config_flow.RADIO_TYPES,
-        {radio_type: {CONTROLLER: app_ctrl_cls, "radio_description": "radio"}},
-    ):
     result = await hass.config_entries.flow.async_configure(
-            result["flow_id"],
+        result["flow_id"], user_input={zigpy.config.CONF_DEVICE_PATH: "/dev/ttyUSB33"},
-            user_input={zigpy.config.CONF_DEVICE_PATH: "/dev/ttyUSB33"},
     )
 
     assert result["type"] == "create_entry"
@@ -187,7 +193,8 @@ async def test_user_port_config(hass, radio_type, orig_ctrl_cls):
         result["data"][zigpy.config.CONF_DEVICE][zigpy.config.CONF_DEVICE_PATH]
         == "/dev/ttyUSB33"
     )
-        assert result["data"][CONF_RADIO_TYPE] == radio_type
+    assert result["data"][CONF_RADIO_TYPE] == "ezsp"
+    assert probe_mock.await_count == 1
 
 
 def test_get_serial_by_id_no_dir():