Add hardware Zigbee flow strategy (home-assistant#153190)

Author: MartinHjelmare

homeassistant/components/homeassistant_hardware/firmware_config_flow.py

@@ -61,6 +61,13 @@ class PickedFirmwareType(StrEnum):
     ZIGBEE = "zigbee"
 
 
+class ZigbeeFlowStrategy(StrEnum):
+    """Zigbee setup strategies that can be picked."""
+
+    ADVANCED = "advanced"
+    RECOMMENDED = "recommended"
+
+
 class ZigbeeIntegration(StrEnum):
     """Zigbee integrations that can be picked."""
 
@@ -73,6 +80,7 @@ class BaseFirmwareInstallFlow(ConfigEntryBaseFlow, ABC):
 
     ZIGBEE_BAUDRATE = 115200  # Default, subclasses may override
     _picked_firmware_type: PickedFirmwareType
+    _zigbee_flow_strategy: ZigbeeFlowStrategy = ZigbeeFlowStrategy.RECOMMENDED
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Instantiate base flow."""
@@ -395,12 +403,14 @@ async def async_step_zigbee_intent_recommended(
     ) -> ConfigFlowResult:
         """Select recommended installation type."""
         self._zigbee_integration = ZigbeeIntegration.ZHA
+        self._zigbee_flow_strategy = ZigbeeFlowStrategy.RECOMMENDED
         return await self._async_continue_picked_firmware()
 
     async def async_step_zigbee_intent_custom(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Select custom installation type."""
+        self._zigbee_flow_strategy = ZigbeeFlowStrategy.ADVANCED
         return await self.async_step_zigbee_integration()
 
     async def async_step_zigbee_integration(
@@ -521,6 +531,7 @@ async def async_step_continue_zigbee(
                     "flow_control": "hardware",
                 },
                 "radio_type": "ezsp",
+                "flow_strategy": self._zigbee_flow_strategy,
             },
         )
         return self._continue_zha_flow(result)

homeassistant/components/zha/config_flow.py

@@ -20,6 +20,9 @@
 from homeassistant.components.file_upload import process_uploaded_file
 from homeassistant.components.hassio import AddonError, AddonState
 from homeassistant.components.homeassistant_hardware import silabs_multiprotocol_addon
+from homeassistant.components.homeassistant_hardware.firmware_config_flow import (
+    ZigbeeFlowStrategy,
+)
 from homeassistant.components.homeassistant_yellow import hardware as yellow_hardware
 from homeassistant.config_entries import (
     SOURCE_IGNORE,
@@ -163,6 +166,7 @@ async def list_serial_ports(hass: HomeAssistant) -> list[ListPortInfo]:
 class BaseZhaFlow(ConfigEntryBaseFlow):
     """Mixin for common ZHA flow steps and forms."""
 
+    _flow_strategy: ZigbeeFlowStrategy | None = None
     _hass: HomeAssistant
     _title: str
 
@@ -373,6 +377,12 @@ async def async_step_choose_setup_strategy(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Choose how to set up the integration from scratch."""
+        if self._flow_strategy == ZigbeeFlowStrategy.RECOMMENDED:
+            # Fast path: automatically form a new network
+            return await self.async_step_setup_strategy_recommended()
+        if self._flow_strategy == ZigbeeFlowStrategy.ADVANCED:
+            # Advanced path: let the user choose
+            return await self.async_step_setup_strategy_advanced()
 
         # Allow onboarding for new users to just create a new network automatically
         if (
@@ -406,6 +416,12 @@ async def async_step_choose_migration_strategy(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Choose how to deal with the current radio's settings during migration."""
+        if self._flow_strategy == ZigbeeFlowStrategy.RECOMMENDED:
+            # Fast path: automatically migrate everything
+            return await self.async_step_migration_strategy_recommended()
+        if self._flow_strategy == ZigbeeFlowStrategy.ADVANCED:
+            # Advanced path: let the user choose
+            return await self.async_step_migration_strategy_advanced()
         return self.async_show_menu(
             step_id="choose_migration_strategy",
             menu_options=[
@@ -867,6 +883,7 @@ async def async_step_hardware(
         radio_type = self._radio_mgr.parse_radio_type(discovery_data["radio_type"])
         device_settings = discovery_data["port"]
         device_path = device_settings[CONF_DEVICE_PATH]
+        self._flow_strategy = discovery_data.get("flow_strategy")
 
         await self._set_unique_id_and_update_ignored_flow(
             unique_id=f"{name}_{radio_type.name}_{device_path}",

homeassistant/components/zha/radio_manager.py

@@ -28,6 +28,9 @@
 
 from homeassistant import config_entries
 from homeassistant.components import usb
+from homeassistant.components.homeassistant_hardware.firmware_config_flow import (
+    ZigbeeFlowStrategy,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.service_info.usb import UsbServiceInfo
@@ -74,6 +77,7 @@
         vol.Required("name"): str,
         vol.Required("port"): DEVICE_SCHEMA,
         vol.Required("radio_type"): str,
+        vol.Optional("flow_strategy"): vol.All(str, vol.Coerce(ZigbeeFlowStrategy)),
     }
 )
 

tests/components/homeassistant_hardware/test_config_flow.py

@@ -364,8 +364,8 @@ async def consume_progress_flow(
     return result
 
 
-async def test_config_flow_recommended(hass: HomeAssistant) -> None:
-    """Test the config flow with recommended installation type for Zigbee."""
+async def test_config_flow_zigbee_recommended(hass: HomeAssistant) -> None:
+    """Test flow with recommended Zigbee installation type."""
     init_result = await hass.config_entries.flow.async_init(
         TEST_DOMAIN, context={"source": "hardware"}
     )
@@ -418,37 +418,28 @@ async def test_config_flow_recommended(hass: HomeAssistant) -> None:
     assert zha_flow["context"]["source"] == "hardware"
     assert zha_flow["step_id"] == "confirm"
 
+    progress_zha_flows = hass.config_entries.flow._async_progress_by_handler(
+        handler="zha",
+        match_context=None,
+    )
 
-@pytest.mark.parametrize(
-    ("zigbee_integration", "zha_flows"),
-    [
-        (
-            "zigbee_integration_zha",
-            [
-                {
-                    "context": {
-                        "confirm_only": True,
-                        "source": "hardware",
-                        "title_placeholders": {
-                            "name": "Some Hardware Name",
-                        },
-                        "unique_id": "Some Hardware Name_ezsp_/dev/SomeDevice123",
-                    },
-                    "flow_id": ANY,
-                    "handler": "zha",
-                    "step_id": "confirm",
-                }
-            ],
-        ),
-        ("zigbee_integration_other", []),
-    ],
-)
-async def test_config_flow_zigbee_custom(
-    hass: HomeAssistant,
-    zigbee_integration: str,
-    zha_flows: list[ConfigFlowResult],
-) -> None:
-    """Test the config flow with custom installation type selected for Zigbee."""
+    assert len(progress_zha_flows) == 1
+
+    progress_zha_flow = progress_zha_flows[0]
+    assert progress_zha_flow.init_data == {
+        "name": "Some Hardware Name",
+        "port": {
+            "path": "/dev/SomeDevice123",
+            "baudrate": 115200,
+            "flow_control": "hardware",
+        },
+        "radio_type": "ezsp",
+        "flow_strategy": "recommended",
+    }
+
+
+async def test_config_flow_zigbee_custom_zha(hass: HomeAssistant) -> None:
+    """Test flow with custom Zigbee installation type and ZHA selected."""
     init_result = await hass.config_entries.flow.async_init(
         TEST_DOMAIN, context={"source": "hardware"}
     )
@@ -479,7 +470,7 @@ async def test_config_flow_zigbee_custom(
 
         pick_result = await hass.config_entries.flow.async_configure(
             pick_result["flow_id"],
-            user_input={"next_step_id": zigbee_integration},
+            user_input={"next_step_id": "zigbee_integration_zha"},
         )
 
         assert pick_result["type"] is FlowResultType.SHOW_PROGRESS
@@ -503,7 +494,98 @@ async def test_config_flow_zigbee_custom(
 
     # Ensure a ZHA discovery flow has been created
     flows = hass.config_entries.flow.async_progress()
-    assert flows == zha_flows
+    assert flows == [
+        {
+            "context": {
+                "confirm_only": True,
+                "source": "hardware",
+                "title_placeholders": {
+                    "name": "Some Hardware Name",
+                },
+                "unique_id": "Some Hardware Name_ezsp_/dev/SomeDevice123",
+            },
+            "flow_id": ANY,
+            "handler": "zha",
+            "step_id": "confirm",
+        }
+    ]
+
+    progress_zha_flows = hass.config_entries.flow._async_progress_by_handler(
+        handler="zha",
+        match_context=None,
+    )
+
+    assert len(progress_zha_flows) == 1
+
+    progress_zha_flow = progress_zha_flows[0]
+    assert progress_zha_flow.init_data == {
+        "name": "Some Hardware Name",
+        "port": {
+            "path": "/dev/SomeDevice123",
+            "baudrate": 115200,
+            "flow_control": "hardware",
+        },
+        "radio_type": "ezsp",
+        "flow_strategy": "advanced",
+    }
+
+
+async def test_config_flow_zigbee_custom_other(hass: HomeAssistant) -> None:
+    """Test flow with custom Zigbee installation type and Other selected."""
+    init_result = await hass.config_entries.flow.async_init(
+        TEST_DOMAIN, context={"source": "hardware"}
+    )
+
+    assert init_result["type"] is FlowResultType.MENU
+    assert init_result["step_id"] == "pick_firmware"
+
+    with mock_firmware_info(
+        probe_app_type=ApplicationType.SPINEL,
+        flash_app_type=ApplicationType.EZSP,
+    ):
+        # Pick the menu option: we are flashing the firmware
+        pick_result = await hass.config_entries.flow.async_configure(
+            init_result["flow_id"],
+            user_input={"next_step_id": STEP_PICK_FIRMWARE_ZIGBEE},
+        )
+
+        assert pick_result["type"] is FlowResultType.MENU
+        assert pick_result["step_id"] == "zigbee_installation_type"
+
+        pick_result = await hass.config_entries.flow.async_configure(
+            pick_result["flow_id"],
+            user_input={"next_step_id": "zigbee_intent_custom"},
+        )
+
+        assert pick_result["type"] is FlowResultType.MENU
+        assert pick_result["step_id"] == "zigbee_integration"
+
+        pick_result = await hass.config_entries.flow.async_configure(
+            pick_result["flow_id"],
+            user_input={"next_step_id": "zigbee_integration_other"},
+        )
+
+        assert pick_result["type"] is FlowResultType.SHOW_PROGRESS
+        assert pick_result["progress_action"] == "install_firmware"
+        assert pick_result["step_id"] == "install_zigbee_firmware"
+
+        create_result = await consume_progress_flow(
+            hass,
+            flow_id=pick_result["flow_id"],
+            valid_step_ids=("install_zigbee_firmware",),
+        )
+
+        assert create_result["type"] is FlowResultType.CREATE_ENTRY
+
+    config_entry = create_result["result"]
+    assert config_entry.data == {
+        "firmware": "ezsp",
+        "device": TEST_DEVICE,
+        "hardware": TEST_HARDWARE_NAME,
+    }
+
+    flows = hass.config_entries.flow.async_progress()
+    assert flows == []
 
 
 async def test_config_flow_firmware_index_download_fails_but_not_required(