Simplify WLED config flow, use device name for config entry (#63377)

homeassistant/components/wled/__init__.py

@@ -23,15 +23,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up WLED from a config entry."""
     coordinator = WLEDDataUpdateCoordinator(hass, entry=entry)
     await coordinator.async_config_entry_first_refresh()
-
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = coordinator
-
-    # For backwards compat, set unique ID
-    if entry.unique_id is None:
-        hass.config_entries.async_update_entry(
-            entry, unique_id=coordinator.data.info.mac_address
-        )
 
     # Set up all platforms for this device/entry.
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
@@ -44,8 +36,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload WLED config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-    if unload_ok:
         coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
 
         # Ensure disconnected and cleanup stop sub

homeassistant/components/wled/config_flow.py

@@ -4,16 +4,11 @@ from __future__ import annotations
 from typing import Any
 
 import voluptuous as vol
-from wled import WLED, WLEDConnectionError
+from wled import WLED, Device, WLEDConnectionError
 
 from homeassistant.components import zeroconf
-from homeassistant.config_entries import (
+from homeassistant.config_entries import ConfigEntry, ConfigFlow, OptionsFlow
-    SOURCE_ZEROCONF,
+from homeassistant.const import CONF_HOST, CONF_MAC
-    ConfigEntry,
-    ConfigFlow,
-    OptionsFlow,
-)
-from homeassistant.const import CONF_HOST, CONF_MAC, CONF_NAME
 from homeassistant.core import callback
 from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
@@ -25,6 +20,8 @@ class WLEDFlowHandler(ConfigFlow, domain=DOMAIN):
     """Handle a WLED config flow."""
 
     VERSION = 1
+    discovered_host: str
+    discovered_device: Device
 
     @staticmethod
     @callback
@@ -36,98 +33,83 @@ class WLEDFlowHandler(ConfigFlow, domain=DOMAIN):
         self, user_input: dict[str, Any] | None = None
     ) -> FlowResult:
         """Handle a flow initiated by the user."""
-        return await self._handle_config_flow(user_input)
+        errors = {}
 
-    async def async_step_zeroconf(
+        if user_input is not None:
-        self, discovery_info: zeroconf.ZeroconfServiceInfo
-    ) -> FlowResult:
-        """Handle zeroconf discovery."""
-
-        # Hostname is format: wled-livingroom.local.
-        host = discovery_info.hostname.rstrip(".")
-        name, _ = host.rsplit(".")
-
-        self.context.update(
-            {
-                CONF_HOST: discovery_info.host,
-                CONF_NAME: name,
-                CONF_MAC: discovery_info.properties.get(CONF_MAC),
-                "title_placeholders": {"name": name},
-            }
-        )
-
-        # Prepare configuration flow
-        return await self._handle_config_flow({}, True)
-
-    async def async_step_zeroconf_confirm(
-        self, user_input: dict[str, Any] | None = None
-    ) -> FlowResult:
-        """Handle a flow initiated by zeroconf."""
-        return await self._handle_config_flow(user_input)
-
-    async def _handle_config_flow(
-        self, user_input: dict[str, Any] | None = None, prepare: bool = False
-    ) -> FlowResult:
-        """Config flow handler for WLED."""
-        source = self.context.get("source")
-
-        # Request user input, unless we are preparing discovery flow
-        if user_input is None and not prepare:
-            if source == SOURCE_ZEROCONF:
-                return self._show_confirm_dialog()
-            return self._show_setup_form()
-
-        # if prepare is True, user_input can not be None.
-        assert user_input is not None
-
-        if source == SOURCE_ZEROCONF:
-            user_input[CONF_HOST] = self.context.get(CONF_HOST)
-            user_input[CONF_MAC] = self.context.get(CONF_MAC)
-
-        if user_input.get(CONF_MAC) is None or not prepare:
-            session = async_get_clientsession(self.hass)
-            wled = WLED(user_input[CONF_HOST], session=session)
             try:
-                device = await wled.update()
+                device = await self._async_get_device(user_input[CONF_HOST])
             except WLEDConnectionError:
-                if source == SOURCE_ZEROCONF:
+                errors["base"] = "cannot_connect"
-                    return self.async_abort(reason="cannot_connect")
+            else:
-                return self._show_setup_form({"base": "cannot_connect"})
+                await self.async_set_unique_id(device.info.mac_address)
-            user_input[CONF_MAC] = device.info.mac_address
+                self._abort_if_unique_id_configured(
+                    updates={CONF_HOST: user_input[CONF_HOST]}
+                )
+                return self.async_create_entry(
+                    title=device.info.name,
+                    data={
+                        CONF_HOST: user_input[CONF_HOST],
+                    },
+                )
+        else:
+            user_input = {}
 
-        # Check if already configured
-        await self.async_set_unique_id(user_input[CONF_MAC])
-        self._abort_if_unique_id_configured(updates={CONF_HOST: user_input[CONF_HOST]})
-
-        title = user_input[CONF_HOST]
-        if source == SOURCE_ZEROCONF:
-            title = self.context.get(CONF_NAME)
-
-        if prepare:
-            return await self.async_step_zeroconf_confirm()
-
-        return self.async_create_entry(
-            title=title,
-            data={CONF_HOST: user_input[CONF_HOST], CONF_MAC: user_input[CONF_MAC]},
-        )
-
-    def _show_setup_form(self, errors: dict | None = None) -> FlowResult:
-        """Show the setup form to the user."""
         return self.async_show_form(
             step_id="user",
             data_schema=vol.Schema({vol.Required(CONF_HOST): str}),
             errors=errors or {},
         )
 
-    def _show_confirm_dialog(self, errors: dict | None = None) -> FlowResult:
+    async def async_step_zeroconf(
-        """Show the confirm dialog to the user."""
+        self, discovery_info: zeroconf.ZeroconfServiceInfo
-        name = self.context.get(CONF_NAME)
+    ) -> FlowResult:
+        """Handle zeroconf discovery."""
+        # Abort quick if the mac address is provided by discovery info
+        if mac := discovery_info.properties.get(CONF_MAC):
+            await self.async_set_unique_id(mac)
+            self._abort_if_unique_id_configured(
+                updates={CONF_HOST: discovery_info.host}
+            )
+
+        self.discovered_host = discovery_info.host
+        try:
+            self.discovered_device = await self._async_get_device(discovery_info.host)
+        except WLEDConnectionError:
+            return self.async_abort(reason="cannot_connect")
+
+        await self.async_set_unique_id(self.discovered_device.info.mac_address)
+        self._abort_if_unique_id_configured(updates={CONF_HOST: discovery_info.host})
+
+        self.context.update(
+            {
+                "title_placeholders": {"name": self.discovered_device.info.name},
+            }
+        )
+        return await self.async_step_zeroconf_confirm()
+
+    async def async_step_zeroconf_confirm(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Handle a flow initiated by zeroconf."""
+        if user_input is not None:
+            return self.async_create_entry(
+                title=self.discovered_device.info.name,
+                data={
+                    CONF_HOST: self.discovered_host,
+                },
+            )
+
         return self.async_show_form(
             step_id="zeroconf_confirm",
-            description_placeholders={"name": name},
+            description_placeholders={"name": self.discovered_device.info.name},
-            errors=errors or {},
         )
 
+    async def _async_get_device(self, host: str) -> Device:
+        """Get device information from WLED device."""
+        session = async_get_clientsession(self.hass)
+        wled = WLED(host, session=session)
+        return await wled.update()
+
 
 class WLEDOptionsFlowHandler(OptionsFlow):
     """Handle WLED options."""

tests/components/wled/conftest.py

@@ -7,7 +7,7 @@ import pytest
 from wled import Device as WLEDDevice
 
 from homeassistant.components.wled.const import DOMAIN
-from homeassistant.const import CONF_HOST, CONF_MAC
+from homeassistant.const import CONF_HOST
 from homeassistant.core import HomeAssistant
 
 from tests.common import MockConfigEntry, load_fixture
@@ -19,7 +19,8 @@ def mock_config_entry() -> MockConfigEntry:
     """Return the default mocked config entry."""
     return MockConfigEntry(
         domain=DOMAIN,
-        data={CONF_HOST: "192.168.1.123", CONF_MAC: "aabbccddeeff"},
+        data={CONF_HOST: "192.168.1.123"},
+        unique_id="aabbccddeeff",
     )
 
 

tests/components/wled/test_config_flow.py

@@ -34,11 +34,12 @@ async def test_full_user_flow_implementation(
         result["flow_id"], user_input={CONF_HOST: "192.168.1.123"}
     )
 
-    assert result.get("title") == "192.168.1.123"
+    assert result.get("title") == "WLED RGB Light"
     assert result.get("type") == RESULT_TYPE_CREATE_ENTRY
     assert "data" in result
     assert result["data"][CONF_HOST] == "192.168.1.123"
-    assert result["data"][CONF_MAC] == "aabbccddeeff"
+    assert "result" in result
+    assert result["result"].unique_id == "aabbccddeeff"
 
 
 async def test_full_zeroconf_flow_implementation(
@@ -53,7 +54,7 @@ async def test_full_zeroconf_flow_implementation(
             hostname="example.local.",
             name="mock_name",
             port=None,
-            properties={},
+            properties={CONF_MAC: "aabbccddeeff"},
             type="mock_type",
         ),
     )
@@ -61,26 +62,22 @@ async def test_full_zeroconf_flow_implementation(
     flows = hass.config_entries.flow.async_progress()
     assert len(flows) == 1
 
-    assert result.get("description_placeholders") == {CONF_NAME: "example"}
+    assert result.get("description_placeholders") == {CONF_NAME: "WLED RGB Light"}
     assert result.get("step_id") == "zeroconf_confirm"
     assert result.get("type") == RESULT_TYPE_FORM
     assert "flow_id" in result
 
-    flow = flows[0]
-    assert "context" in flow
-    assert flow["context"][CONF_HOST] == "192.168.1.123"
-    assert flow["context"][CONF_NAME] == "example"
-
     result2 = await hass.config_entries.flow.async_configure(
         result["flow_id"], user_input={}
     )
 
-    assert result2.get("title") == "example"
+    assert result2.get("title") == "WLED RGB Light"
     assert result2.get("type") == RESULT_TYPE_CREATE_ENTRY
 
     assert "data" in result2
     assert result2["data"][CONF_HOST] == "192.168.1.123"
-    assert result2["data"][CONF_MAC] == "aabbccddeeff"
+    assert "result" in result2
+    assert result2["result"].unique_id == "aabbccddeeff"
 
 
 async def test_connection_error(
@@ -113,34 +110,7 @@ async def test_zeroconf_connection_error(
             hostname="example.local.",
             name="mock_name",
             port=None,
-            properties={},
+            properties={CONF_MAC: "aabbccddeeff"},
-            type="mock_type",
-        ),
-    )
-
-    assert result.get("type") == RESULT_TYPE_ABORT
-    assert result.get("reason") == "cannot_connect"
-
-
-async def test_zeroconf_confirm_connection_error(
-    hass: HomeAssistant, mock_wled_config_flow: MagicMock
-) -> None:
-    """Test we abort zeroconf flow on WLED connection error."""
-    mock_wled_config_flow.update.side_effect = WLEDConnectionError
-
-    result = await hass.config_entries.flow.async_init(
-        DOMAIN,
-        context={
-            "source": SOURCE_ZEROCONF,
-            CONF_HOST: "example.com",
-            CONF_NAME: "test",
-        },
-        data=zeroconf.ZeroconfServiceInfo(
-            host="192.168.1.123",
-            hostname="example.com.",
-            name="mock_name",
-            port=None,
-            properties={},
             type="mock_type",
         ),
     )
@@ -151,10 +121,11 @@ async def test_zeroconf_confirm_connection_error(
 
 async def test_user_device_exists_abort(
     hass: HomeAssistant,
-    init_integration: MagicMock,
+    mock_config_entry: MockConfigEntry,
     mock_wled_config_flow: MagicMock,
 ) -> None:
     """Test we abort zeroconf flow if WLED device already configured."""
+    mock_config_entry.add_to_hass(hass)
     result = await hass.config_entries.flow.async_init(
         DOMAIN,
         context={"source": SOURCE_USER},
@@ -165,12 +136,13 @@ async def test_user_device_exists_abort(
     assert result.get("reason") == "already_configured"
 
 
-async def test_zeroconf_device_exists_abort(
+async def test_zeroconf_without_mac_device_exists_abort(
     hass: HomeAssistant,
-    init_integration: MagicMock,
+    mock_config_entry: MockConfigEntry,
     mock_wled_config_flow: MagicMock,
 ) -> None:
     """Test we abort zeroconf flow if WLED device already configured."""
+    mock_config_entry.add_to_hass(hass)
     result = await hass.config_entries.flow.async_init(
         DOMAIN,
         context={"source": SOURCE_ZEROCONF},
@@ -190,10 +162,11 @@ async def test_zeroconf_device_exists_abort(
 
 async def test_zeroconf_with_mac_device_exists_abort(
     hass: HomeAssistant,
-    init_integration: MockConfigEntry,
+    mock_config_entry: MockConfigEntry,
     mock_wled_config_flow: MagicMock,
 ) -> None:
     """Test we abort zeroconf flow if WLED device already configured."""
+    mock_config_entry.add_to_hass(hass)
     result = await hass.config_entries.flow.async_init(
         DOMAIN,
         context={"source": SOURCE_ZEROCONF},