ZHA Yellow config flow fixes (#77603)

homeassistant/components/zha/api.py

@@ -50,6 +50,7 @@ from .core.const import (
     DATA_ZHA,
     DATA_ZHA_GATEWAY,
     DOMAIN,
+    EZSP_OVERWRITE_EUI64,
     GROUP_ID,
     GROUP_IDS,
     GROUP_NAME,
@@ -1140,7 +1141,7 @@ async def websocket_restore_network_backup(
 
     if msg["ezsp_force_write_eui64"]:
         backup.network_info.stack_specific.setdefault("ezsp", {})[
-            "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it"
+            EZSP_OVERWRITE_EUI64
         ] = True
 
     # This can take 30-40s

homeassistant/components/zha/config_flow.py

@@ -35,6 +35,7 @@ from .core.const import (
     DATA_ZHA_CONFIG,
     DEFAULT_DATABASE_NAME,
     DOMAIN,
+    EZSP_OVERWRITE_EUI64,
     RadioType,
 )
 
@@ -91,9 +92,7 @@ def _allow_overwrite_ezsp_ieee(
 ) -> zigpy.backups.NetworkBackup:
     """Return a new backup with the flag to allow overwriting the EZSP EUI64."""
     new_stack_specific = copy.deepcopy(backup.network_info.stack_specific)
-    new_stack_specific.setdefault("ezsp", {})[
+    new_stack_specific.setdefault("ezsp", {})[EZSP_OVERWRITE_EUI64] = True
-        "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it"
-    ] = True
 
     return backup.replace(
         network_info=backup.network_info.replace(stack_specific=new_stack_specific)
@@ -108,9 +107,7 @@ def _prevent_overwrite_ezsp_ieee(
         return backup
 
     new_stack_specific = copy.deepcopy(backup.network_info.stack_specific)
-    new_stack_specific.setdefault("ezsp", {}).pop(
+    new_stack_specific.setdefault("ezsp", {}).pop(EZSP_OVERWRITE_EUI64, None)
-        "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it", None
-    )
 
     return backup.replace(
         network_info=backup.network_info.replace(stack_specific=new_stack_specific)
@@ -664,10 +661,12 @@ class ZhaConfigFlowHandler(BaseZhaFlow, config_entries.ConfigFlow, domain=DOMAIN
         """Handle hardware flow."""
         if self._async_current_entries():
             return self.async_abort(reason="single_instance_allowed")
+
         if not data:
             return self.async_abort(reason="invalid_hardware_data")
         if data.get("radio_type") != "efr32":
             return self.async_abort(reason="invalid_hardware_data")
+
         self._radio_type = RadioType.ezsp
 
         schema = {
@@ -689,23 +688,10 @@ class ZhaConfigFlowHandler(BaseZhaFlow, config_entries.ConfigFlow, domain=DOMAIN
             return self.async_abort(reason="invalid_hardware_data")
 
         self._title = data.get("name", data["port"]["path"])
-        self._device_path = device_settings.pop(CONF_DEVICE_PATH)
+        self._device_path = device_settings[CONF_DEVICE_PATH]
         self._device_settings = device_settings
 
-        self._set_confirm_only()
+        return await self.async_step_choose_formation_strategy()
-        return await self.async_step_confirm_hardware()
-
-    async def async_step_confirm_hardware(
-        self, user_input: dict[str, Any] | None = None
-    ) -> FlowResult:
-        """Confirm a hardware discovery."""
-        if user_input is not None or not onboarding.async_is_onboarded(self.hass):
-            return await self._async_create_radio_entity()
-
-        return self.async_show_form(
-            step_id="confirm_hardware",
-            description_placeholders={CONF_NAME: self._title},
-        )
 
 
 class ZhaOptionsFlowHandler(BaseZhaFlow, config_entries.OptionsFlow):

homeassistant/components/zha/core/const.py

@@ -412,3 +412,7 @@ class Strobe(t.enum8):
 
 STARTUP_FAILURE_DELAY_S = 3
 STARTUP_RETRIES = 3
+
+EZSP_OVERWRITE_EUI64 = (
+    "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it"
+)

tests/components/homeassistant_yellow/conftest.py

@@ -1,13 +1,28 @@
 """Test fixtures for the Home Assistant Yellow integration."""
-from unittest.mock import patch
+from collections.abc import Generator
+from typing import Any
+from unittest.mock import MagicMock, patch
 
 import pytest
 
 
 @pytest.fixture(autouse=True)
-def mock_zha():
+def mock_zha_config_flow_setup() -> Generator[None, None, None]:
-    """Mock the zha integration."""
+    """Mock the radio connection and probing of the ZHA config flow."""
+
+    def mock_probe(config: dict[str, Any]) -> None:
+        # The radio probing will return the correct baudrate
+        return {**config, "baudrate": 115200}
+
+    mock_connect_app = MagicMock()
+    mock_connect_app.__aenter__.return_value.backups.backups = []
+
     with patch(
+        "bellows.zigbee.application.ControllerApplication.probe", side_effect=mock_probe
+    ), patch(
+        "homeassistant.components.zha.config_flow.BaseZhaFlow._connect_zigpy_app",
+        return_value=mock_connect_app,
+    ), patch(
         "homeassistant.components.zha.async_setup_entry",
         return_value=True,
     ):

tests/components/homeassistant_yellow/test_init.py

@@ -3,6 +3,7 @@ from unittest.mock import patch
 
 import pytest
 
+from homeassistant.components import zha
 from homeassistant.components.homeassistant_yellow.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
@@ -37,8 +38,20 @@ async def test_setup_entry(
         await hass.async_block_till_done()
         assert len(mock_get_os_info.mock_calls) == 1
 
-    assert len(hass.config_entries.async_entries("zha")) == num_entries
+    # Finish setting up ZHA
+    if num_entries > 0:
+        zha_flows = hass.config_entries.flow.async_progress_by_handler("zha")
+        assert len(zha_flows) == 1
+        assert zha_flows[0]["step_id"] == "choose_formation_strategy"
+
+        await hass.config_entries.flow.async_configure(
+            zha_flows[0]["flow_id"],
+            user_input={"next_step_id": zha.config_flow.FORMATION_REUSE_SETTINGS},
+        )
+        await hass.async_block_till_done()
+
     assert len(hass.config_entries.flow.async_progress_by_handler("zha")) == num_flows
+    assert len(hass.config_entries.async_entries("zha")) == num_entries
 
 
 async def test_setup_zha(hass: HomeAssistant) -> None:
@@ -63,6 +76,17 @@ async def test_setup_zha(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
         assert len(mock_get_os_info.mock_calls) == 1
 
+    # Finish setting up ZHA
+    zha_flows = hass.config_entries.flow.async_progress_by_handler("zha")
+    assert len(zha_flows) == 1
+    assert zha_flows[0]["step_id"] == "choose_formation_strategy"
+
+    await hass.config_entries.flow.async_configure(
+        zha_flows[0]["flow_id"],
+        user_input={"next_step_id": zha.config_flow.FORMATION_REUSE_SETTINGS},
+    )
+    await hass.async_block_till_done()
+
     config_entry = hass.config_entries.async_entries("zha")[0]
     assert config_entry.data == {
         "device": {

tests/components/zha/test_api.py

@@ -34,6 +34,7 @@ from homeassistant.components.zha.core.const import (
     CLUSTER_TYPE_IN,
     DATA_ZHA,
     DATA_ZHA_GATEWAY,
+    EZSP_OVERWRITE_EUI64,
     GROUP_ID,
     GROUP_IDS,
     GROUP_NAME,
@@ -709,11 +710,7 @@ async def test_restore_network_backup_force_write_eui64(app_controller, zha_clie
     p.assert_called_once_with(
         backup.replace(
             network_info=backup.network_info.replace(
-                stack_specific={
+                stack_specific={"ezsp": {EZSP_OVERWRITE_EUI64: True}}
-                    "ezsp": {
-                        "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it": True
-                    }
-                }
             )
         )
     )

tests/components/zha/test_config_flow.py

@@ -21,6 +21,7 @@ from homeassistant.components.zha.core.const import (
     CONF_FLOWCONTROL,
     CONF_RADIO_TYPE,
     DOMAIN,
+    EZSP_OVERWRITE_EUI64,
     RadioType,
 )
 from homeassistant.config_entries import (
@@ -857,8 +858,9 @@ async def test_migration_ti_cc_to_znp(old_type, new_type, hass, config_entry):
     assert config_entry.data[CONF_RADIO_TYPE] == new_type
 
 
+@pytest.mark.parametrize("onboarded", [True, False])
 @patch("homeassistant.components.zha.async_setup_entry", AsyncMock(return_value=True))
-async def test_hardware_not_onboarded(hass):
+async def test_hardware(onboarded, hass):
     """Test hardware flow."""
     data = {
         "name": "Yellow",
@@ -870,52 +872,23 @@ async def test_hardware_not_onboarded(hass):
         },
     }
     with patch(
-        "homeassistant.components.onboarding.async_is_onboarded", return_value=False
+        "homeassistant.components.onboarding.async_is_onboarded", return_value=onboarded
     ):
-        result = await hass.config_entries.flow.async_init(
+        result1 = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": "hardware"}, data=data
         )
 
-    assert result["type"] == FlowResultType.CREATE_ENTRY
+    assert result1["type"] == FlowResultType.MENU
-    assert result["title"] == "Yellow"
+    assert result1["step_id"] == "choose_formation_strategy"
-    assert result["data"] == {
-        CONF_DEVICE: {
-            CONF_BAUDRATE: 115200,
-            CONF_FLOWCONTROL: "hardware",
-            CONF_DEVICE_PATH: "/dev/ttyAMA1",
-        },
-        CONF_RADIO_TYPE: "ezsp",
-    }
 
-
+    result2 = await hass.config_entries.flow.async_configure(
-@patch("homeassistant.components.zha.async_setup_entry", AsyncMock(return_value=True))
+        result1["flow_id"],
-async def test_hardware_onboarded(hass):
+        user_input={"next_step_id": config_flow.FORMATION_REUSE_SETTINGS},
-    """Test hardware flow."""
-    data = {
-        "radio_type": "efr32",
-        "port": {
-            "path": "/dev/ttyAMA1",
-            "baudrate": 115200,
-            "flow_control": "hardware",
-        },
-    }
-    with patch(
-        "homeassistant.components.onboarding.async_is_onboarded", return_value=True
-    ):
-        result = await hass.config_entries.flow.async_init(
-            DOMAIN, context={"source": "hardware"}, data=data
-        )
-
-    assert result["type"] == FlowResultType.FORM
-    assert result["step_id"] == "confirm_hardware"
-
-    result = await hass.config_entries.flow.async_configure(
-        result["flow_id"], user_input={}
     )
+    await hass.async_block_till_done()
 
-    assert result["type"] == FlowResultType.CREATE_ENTRY
+    assert result2["title"] == "Yellow"
-    assert result["title"] == "/dev/ttyAMA1"
+    assert result2["data"] == {
-    assert result["data"] == {
         CONF_DEVICE: {
             CONF_BAUDRATE: 115200,
             CONF_FLOWCONTROL: "hardware",
@@ -968,25 +941,18 @@ def test_allow_overwrite_ezsp_ieee():
     new_backup = config_flow._allow_overwrite_ezsp_ieee(backup)
 
     assert backup != new_backup
-    assert (
+    assert new_backup.network_info.stack_specific["ezsp"][EZSP_OVERWRITE_EUI64] is True
-        new_backup.network_info.stack_specific["ezsp"][
-            "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it"
-        ]
-        is True
-    )
 
 
 def test_prevent_overwrite_ezsp_ieee():
     """Test modifying the backup to prevent bellows from overriding the IEEE address."""
     backup = zigpy.backups.NetworkBackup()
-    backup.network_info.stack_specific["ezsp"] = {
+    backup.network_info.stack_specific["ezsp"] = {EZSP_OVERWRITE_EUI64: True}
-        "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it": True
-    }
     new_backup = config_flow._prevent_overwrite_ezsp_ieee(backup)
 
     assert backup != new_backup
     assert not new_backup.network_info.stack_specific.get("ezsp", {}).get(
-        "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it"
+        EZSP_OVERWRITE_EUI64
     )
 
 
@@ -1356,9 +1322,7 @@ async def test_ezsp_restore_without_settings_change_ieee(
     mock_app.state.network_info.network_key.tx_counter += 10000
 
     # Include the overwrite option, just in case someone uploads a backup with it
-    backup.network_info.metadata["ezsp"] = {
+    backup.network_info.metadata["ezsp"] = {EZSP_OVERWRITE_EUI64: True}
-        "i_understand_i_can_update_eui64_only_once_and_i_still_want_to_do_it": True
-    }
 
     result2 = await hass.config_entries.flow.async_configure(
         result["flow_id"],