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

This commit is contained in:
Paulus Schoutsen
2025-10-05 16:33:12 -04:00
committed by GitHub
parent 5d83c82b81
commit 19f990ed31
2 changed files with 249 additions and 6 deletions

View File

@@ -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
def _async_create_entry(self) -> ConfigFlowResult:
async 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:

View File

@@ -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