From cd9096907b5ce9461bfcffeccc5ab4b488875f02 Mon Sep 17 00:00:00 2001 From: Franck Nijhof Date: Tue, 4 Jan 2022 19:59:14 +0100 Subject: [PATCH] Simplify WLED config flow, use device name for config entry (#63377) --- homeassistant/components/wled/__init__.py | 13 +- homeassistant/components/wled/config_flow.py | 154 ++++++++----------- tests/components/wled/conftest.py | 5 +- tests/components/wled/test_config_flow.py | 59 ++----- 4 files changed, 89 insertions(+), 142 deletions(-) diff --git a/homeassistant/components/wled/__init__.py b/homeassistant/components/wled/__init__.py index 659df1baad9..5a25a461fd6 100644 --- a/homeassistant/components/wled/__init__.py +++ b/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, {}) - 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 - ) + hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator # 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: + if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] # Ensure disconnected and cleanup stop sub diff --git a/homeassistant/components/wled/config_flow.py b/homeassistant/components/wled/config_flow.py index 485afef4f6c..bc52467b1ef 100644 --- a/homeassistant/components/wled/config_flow.py +++ b/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 ( - SOURCE_ZEROCONF, - ConfigEntry, - ConfigFlow, - OptionsFlow, -) -from homeassistant.const import CONF_HOST, CONF_MAC, CONF_NAME +from homeassistant.config_entries import ConfigEntry, ConfigFlow, OptionsFlow +from homeassistant.const import CONF_HOST, CONF_MAC 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( - 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) + if user_input is not None: try: - device = await wled.update() + device = await self._async_get_device(user_input[CONF_HOST]) except WLEDConnectionError: - if source == SOURCE_ZEROCONF: - return self.async_abort(reason="cannot_connect") - return self._show_setup_form({"base": "cannot_connect"}) - user_input[CONF_MAC] = device.info.mac_address + errors["base"] = "cannot_connect" + else: + await self.async_set_unique_id(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: - """Show the confirm dialog to the user.""" - name = self.context.get(CONF_NAME) + async def async_step_zeroconf( + self, discovery_info: zeroconf.ZeroconfServiceInfo + ) -> 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}, - errors=errors or {}, + description_placeholders={"name": self.discovered_device.info.name}, ) + 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.""" diff --git a/tests/components/wled/conftest.py b/tests/components/wled/conftest.py index 708bbf46834..595ab3e978d 100644 --- a/tests/components/wled/conftest.py +++ b/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", ) diff --git a/tests/components/wled/test_config_flow.py b/tests/components/wled/test_config_flow.py index 770d6abd2f8..7459fcfc80d 100644 --- a/tests/components/wled/test_config_flow.py +++ b/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={}, - 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={}, + properties={CONF_MAC: "aabbccddeeff"}, 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},