ESPHome to set Z-Wave discovery as next_flow (#153706)

homeassistant/components/esphome/config_flow.py

@@ -22,19 +22,23 @@ import voluptuous as vol
 
 from homeassistant.components import zeroconf
 from homeassistant.config_entries import (
+    SOURCE_ESPHOME,
     SOURCE_IGNORE,
     SOURCE_REAUTH,
     SOURCE_RECONFIGURE,
     ConfigEntry,
     ConfigFlow,
     ConfigFlowResult,
+    FlowType,
     OptionsFlow,
 )
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT
 from homeassistant.core import callback
-from homeassistant.data_entry_flow import AbortFlow
+from homeassistant.data_entry_flow import AbortFlow, FlowResultType
+from homeassistant.helpers import discovery_flow
 from homeassistant.helpers.device_registry import format_mac
 from homeassistant.helpers.service_info.dhcp import DhcpServiceInfo
+from homeassistant.helpers.service_info.esphome import ESPHomeServiceInfo
 from homeassistant.helpers.service_info.hassio import HassioServiceInfo
 from homeassistant.helpers.service_info.mqtt import MqttServiceInfo
 from homeassistant.helpers.service_info.zeroconf import ZeroconfServiceInfo
@@ -75,6 +79,7 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
     def __init__(self) -> None:
         """Initialize flow."""
         self._host: str | None = None
+        self._connected_address: str | None = None
         self.__name: str | None = None
         self._port: int | None = None
         self._password: str | None = None
@@ -498,18 +503,55 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
         await self.hass.config_entries.async_remove(
             self._entry_with_name_conflict.entry_id
         )
-        return self._async_create_entry()
+        return await self._async_create_entry()
 
-    @callback
+    async def _async_create_entry(self) -> ConfigFlowResult:
-    def _async_create_entry(self) -> ConfigFlowResult:
         """Create the config entry."""
         assert self._name is not None
+        assert self._device_info is not None
+
+        # Check if Z-Wave capabilities are present and start discovery flow
+        next_flow_id: str | None = None
+        if self._device_info.zwave_proxy_feature_flags:
+            assert self._connected_address is not None
+            assert self._port is not None
+
+            # Start Z-Wave discovery flow and get the flow ID
+            zwave_result = await self.hass.config_entries.flow.async_init(
+                "zwave_js",
+                context={
+                    "source": SOURCE_ESPHOME,
+                    "discovery_key": discovery_flow.DiscoveryKey(
+                        domain=DOMAIN,
+                        key=self._device_info.mac_address,
+                        version=1,
+                    ),
+                },
+                data=ESPHomeServiceInfo(
+                    name=self._device_info.name,
+                    zwave_home_id=self._device_info.zwave_home_id or None,
+                    ip_address=self._connected_address,
+                    port=self._port,
+                    noise_psk=self._noise_psk,
+                ),
+            )
+            if zwave_result["type"] in (
+                FlowResultType.ABORT,
+                FlowResultType.CREATE_ENTRY,
+            ):
+                _LOGGER.debug(
+                    "Unable to continue created Z-Wave JS config flow: %s", zwave_result
+                )
+            else:
+                next_flow_id = zwave_result["flow_id"]
+
         return self.async_create_entry(
             title=self._name,
             data=self._async_make_config_data(),
             options={
                 CONF_ALLOW_SERVICE_CALLS: DEFAULT_NEW_CONFIG_ALLOW_ALLOW_SERVICE_CALLS,
             },
+            next_flow=(FlowType.CONFIG_FLOW, next_flow_id) if next_flow_id else None,
         )
 
     @callback
@@ -556,7 +598,7 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
             if entry.data.get(CONF_DEVICE_NAME) == self._device_name:
                 self._entry_with_name_conflict = entry
                 return await self.async_step_name_conflict()
-        return self._async_create_entry()
+        return await self._async_create_entry()
 
     async def _async_reauth_validated_connection(self) -> ConfigFlowResult:
         """Handle reauth validated connection."""
@@ -703,6 +745,7 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
         try:
             await cli.connect()
             self._device_info = await cli.device_info()
+            self._connected_address = cli.connected_address
         except InvalidAuthAPIError:
             return ERROR_INVALID_PASSWORD_AUTH
         except RequiresEncryptionAPIError:

tests/components/esphome/test_config_flow.py

@@ -3,7 +3,7 @@
 from ipaddress import ip_address
 import json
 from typing import Any
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from aioesphomeapi import (
     APIClient,
@@ -34,7 +34,9 @@ from homeassistant.config_entries import SOURCE_IGNORE, ConfigFlowResult
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
+from homeassistant.helpers import discovery_flow
 from homeassistant.helpers.service_info.dhcp import DhcpServiceInfo
+from homeassistant.helpers.service_info.esphome import ESPHomeServiceInfo
 from homeassistant.helpers.service_info.hassio import HassioServiceInfo
 from homeassistant.helpers.service_info.mqtt import MqttServiceInfo
 from homeassistant.helpers.service_info.zeroconf import ZeroconfServiceInfo
@@ -2619,3 +2621,201 @@ async def test_discovery_dhcp_no_probe_same_host_port_none(
 
     # Host should remain unchanged
     assert entry.data[CONF_HOST] == "192.168.43.183"
+
+
+@pytest.mark.usefixtures("mock_setup_entry", "mock_zeroconf")
+async def test_user_flow_starts_zwave_discovery(
+    hass: HomeAssistant, mock_client: APIClient
+) -> None:
+    """Test that the user flow starts Z-Wave JS discovery when device has Z-Wave capabilities."""
+    # Mock device with Z-Wave capabilities
+    mock_client.device_info = AsyncMock(
+        return_value=DeviceInfo(
+            uses_password=False,
+            name="test-zwave-device",
+            mac_address="11:22:33:44:55:BB",
+            zwave_proxy_feature_flags=1,
+            zwave_home_id=1234567890,
+        )
+    )
+    mock_client.connected_address = "mock-connected-address"
+
+    # Track flow.async_init calls and async_get calls
+    original_async_init = hass.config_entries.flow.async_init
+    original_async_get = hass.config_entries.flow.async_get
+    flow_init_calls = []
+    zwave_flow_id = "mock-zwave-flow-id"
+
+    async def track_async_init(*args, **kwargs):
+        flow_init_calls.append((args, kwargs))
+        # For the Z-Wave flow, return a mock result with the flow_id
+        if args and args[0] == "zwave_js":
+            return {"flow_id": zwave_flow_id, "type": FlowResultType.FORM}
+        # Otherwise call the original
+        return await original_async_init(*args, **kwargs)
+
+    def mock_async_get(flow_id: str):
+        # Return a mock flow for the Z-Wave flow_id
+        if flow_id == zwave_flow_id:
+            return MagicMock()
+        return original_async_get(flow_id)
+
+    with (
+        patch.object(
+            hass.config_entries.flow, "async_init", side_effect=track_async_init
+        ),
+        patch.object(hass.config_entries.flow, "async_get", side_effect=mock_async_get),
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_USER},
+            data={CONF_HOST: "192.168.1.100", CONF_PORT: 6053},
+        )
+
+    # Verify the entry was created
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert result["title"] == "test-zwave-device"
+    assert result["data"] == {
+        CONF_HOST: "192.168.1.100",
+        CONF_PORT: 6053,
+        CONF_PASSWORD: "",
+        CONF_NOISE_PSK: "",
+        CONF_DEVICE_NAME: "test-zwave-device",
+    }
+
+    # First call is ESPHome flow, second should be Z-Wave flow
+    assert len(flow_init_calls) == 2
+    zwave_call_args, zwave_call_kwargs = flow_init_calls[1]
+    assert zwave_call_args[0] == "zwave_js"
+    assert zwave_call_kwargs["context"] == {
+        "source": config_entries.SOURCE_ESPHOME,
+        "discovery_key": discovery_flow.DiscoveryKey(
+            domain="esphome", key="11:22:33:44:55:BB", version=1
+        ),
+    }
+    assert zwave_call_kwargs["data"] == ESPHomeServiceInfo(
+        name="test-zwave-device",
+        zwave_home_id=1234567890,
+        ip_address="mock-connected-address",
+        port=6053,
+        noise_psk=None,
+    )
+
+    # Verify next_flow was set
+    assert result["next_flow"] == (config_entries.FlowType.CONFIG_FLOW, zwave_flow_id)
+
+
+@pytest.mark.usefixtures("mock_setup_entry", "mock_zeroconf")
+async def test_user_flow_no_zwave_discovery_without_capabilities(
+    hass: HomeAssistant, mock_client: APIClient
+) -> None:
+    """Test that the user flow does not start Z-Wave JS discovery when device has no Z-Wave capabilities."""
+    # Mock device without Z-Wave capabilities
+    mock_client.device_info = AsyncMock(
+        return_value=DeviceInfo(
+            uses_password=False,
+            name="test-regular-device",
+            mac_address="11:22:33:44:55:CC",
+        )
+    )
+
+    # Track flow.async_init calls
+    original_async_init = hass.config_entries.flow.async_init
+    flow_init_calls = []
+
+    async def track_async_init(*args, **kwargs):
+        flow_init_calls.append((args, kwargs))
+        return await original_async_init(*args, **kwargs)
+
+    with patch.object(
+        hass.config_entries.flow, "async_init", side_effect=track_async_init
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_USER},
+            data={CONF_HOST: "192.168.1.101", CONF_PORT: 6053},
+        )
+
+    # Verify the entry was created
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert result["title"] == "test-regular-device"
+
+    # Verify Z-Wave discovery flow was NOT started (only ESPHome flow)
+    assert len(flow_init_calls) == 1
+
+    # Verify next_flow was not set
+    assert "next_flow" not in result
+
+
+@pytest.mark.usefixtures("mock_setup_entry", "mock_zeroconf")
+async def test_user_flow_zwave_discovery_aborts(
+    hass: HomeAssistant, mock_client: APIClient
+) -> None:
+    """Test that the user flow handles Z-Wave discovery abort gracefully."""
+    # Mock device with Z-Wave capabilities
+    mock_client.device_info = AsyncMock(
+        return_value=DeviceInfo(
+            uses_password=False,
+            name="test-zwave-device",
+            mac_address="11:22:33:44:55:DD",
+            zwave_proxy_feature_flags=1,
+            zwave_home_id=9876543210,
+        )
+    )
+    mock_client.connected_address = "192.168.1.102"
+
+    # Track flow.async_init calls
+    original_async_init = hass.config_entries.flow.async_init
+    flow_init_calls = []
+
+    async def track_async_init(*args, **kwargs):
+        flow_init_calls.append((args, kwargs))
+        # For the Z-Wave flow, return an ABORT result
+        if args and args[0] == "zwave_js":
+            return {
+                "type": FlowResultType.ABORT,
+                "reason": "already_configured",
+            }
+        # Otherwise call the original
+        return await original_async_init(*args, **kwargs)
+
+    with patch.object(
+        hass.config_entries.flow, "async_init", side_effect=track_async_init
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_USER},
+            data={CONF_HOST: "192.168.1.102", CONF_PORT: 6053},
+        )
+
+    # Verify the ESPHome entry was still created despite Z-Wave flow aborting
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert result["title"] == "test-zwave-device"
+    assert result["data"] == {
+        CONF_HOST: "192.168.1.102",
+        CONF_PORT: 6053,
+        CONF_PASSWORD: "",
+        CONF_NOISE_PSK: "",
+        CONF_DEVICE_NAME: "test-zwave-device",
+    }
+
+    # Verify Z-Wave discovery flow was attempted
+    assert len(flow_init_calls) == 2
+    zwave_call_args, zwave_call_kwargs = flow_init_calls[1]
+    assert zwave_call_args[0] == "zwave_js"
+    assert zwave_call_kwargs["context"]["source"] == config_entries.SOURCE_ESPHOME
+    assert zwave_call_kwargs["context"]["discovery_key"] == discovery_flow.DiscoveryKey(
+        domain=DOMAIN,
+        key="11:22:33:44:55:DD",
+        version=1,
+    )
+    assert zwave_call_kwargs["data"] == ESPHomeServiceInfo(
+        name="test-zwave-device",
+        zwave_home_id=9876543210,
+        ip_address="192.168.1.102",
+        port=6053,
+        noise_psk=None,
+    )
+
+    # Verify next_flow was NOT set since Z-Wave flow aborted
+    assert "next_flow" not in result