From 9dc158f5e02682fbeb203f80d0ae2be72868d901 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 5 Feb 2022 15:23:31 -0600 Subject: [PATCH] Add support for picking discovered devices to WiZ (#65826) * Add support for picking discovered devices - Also fixes state not being written initially (it was not so obvious since the next coordinator update wrote it) * store it * store it * order * fixes * more cleanups * hints * naming * merge branches --- homeassistant/components/wiz/config_flow.py | 56 +++++++- homeassistant/components/wiz/entity.py | 7 +- homeassistant/components/wiz/light.py | 59 +++----- homeassistant/components/wiz/strings.json | 7 +- homeassistant/components/wiz/switch.py | 5 + .../components/wiz/translations/en.json | 7 +- tests/components/wiz/test_config_flow.py | 135 ++++++++++++++++-- 7 files changed, 223 insertions(+), 53 deletions(-) diff --git a/homeassistant/components/wiz/config_flow.py b/homeassistant/components/wiz/config_flow.py index 34b484f145e..9b753284dc5 100644 --- a/homeassistant/components/wiz/config_flow.py +++ b/homeassistant/components/wiz/config_flow.py @@ -14,11 +14,14 @@ from homeassistant.components import dhcp from homeassistant.const import CONF_HOST from homeassistant.data_entry_flow import FlowResult -from .const import DOMAIN, WIZ_EXCEPTIONS -from .utils import name_from_bulb_type_and_mac +from .const import DEFAULT_NAME, DISCOVER_SCAN_TIMEOUT, DOMAIN, WIZ_EXCEPTIONS +from .discovery import async_discover_devices +from .utils import _short_mac, name_from_bulb_type_and_mac _LOGGER = logging.getLogger(__name__) +CONF_DEVICE = "device" + class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): """Handle a config flow for WiZ.""" @@ -28,6 +31,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): def __init__(self) -> None: """Initialize the config flow.""" self._discovered_device: DiscoveredBulb | None = None + self._discovered_devices: dict[str, DiscoveredBulb] = {} self._name: str | None = None async def async_step_dhcp(self, discovery_info: dhcp.DhcpServiceInfo) -> FlowResult: @@ -85,13 +89,57 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): data_schema=vol.Schema({}), ) + async def async_step_pick_device( + self, user_input: dict[str, Any] | None = None + ) -> FlowResult: + """Handle the step to pick discovered device.""" + if user_input is not None: + device = self._discovered_devices[user_input[CONF_DEVICE]] + await self.async_set_unique_id(device.mac_address, raise_on_progress=False) + bulb = wizlight(device.ip_address) + try: + bulbtype = await bulb.get_bulbtype() + except WIZ_EXCEPTIONS: + return self.async_abort(reason="cannot_connect") + else: + return self.async_create_entry( + title=name_from_bulb_type_and_mac(bulbtype, device.mac_address), + data={CONF_HOST: device.ip_address}, + ) + + current_unique_ids = self._async_current_ids() + current_hosts = { + entry.data[CONF_HOST] + for entry in self._async_current_entries(include_ignore=False) + } + discovered_devices = await async_discover_devices( + self.hass, DISCOVER_SCAN_TIMEOUT + ) + self._discovered_devices = { + device.mac_address: device for device in discovered_devices + } + devices_name = { + mac: f"{DEFAULT_NAME} {_short_mac(mac)} ({device.ip_address})" + for mac, device in self._discovered_devices.items() + if mac not in current_unique_ids and device.ip_address not in current_hosts + } + # Check if there is at least one device + if not devices_name: + return self.async_abort(reason="no_devices_found") + return self.async_show_form( + step_id="pick_device", + data_schema=vol.Schema({vol.Required(CONF_DEVICE): vol.In(devices_name)}), + ) + async def async_step_user( self, user_input: dict[str, Any] | None = None ) -> FlowResult: """Handle a flow initialized by the user.""" errors = {} if user_input is not None: - bulb = wizlight(user_input[CONF_HOST]) + if not (host := user_input[CONF_HOST]): + return await self.async_step_pick_device() + bulb = wizlight(host) try: mac = await bulb.getMac() bulbtype = await bulb.get_bulbtype() @@ -117,6 +165,6 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): return self.async_show_form( step_id="user", - data_schema=vol.Schema({vol.Required(CONF_HOST): str}), + data_schema=vol.Schema({vol.Optional(CONF_HOST, default=""): str}), errors=errors, ) diff --git a/homeassistant/components/wiz/entity.py b/homeassistant/components/wiz/entity.py index 0c4829383d3..de0b0e3f947 100644 --- a/homeassistant/components/wiz/entity.py +++ b/homeassistant/components/wiz/entity.py @@ -33,9 +33,14 @@ class WizToggleEntity(CoordinatorEntity, ToggleEntity): @callback def _handle_coordinator_update(self) -> None: """Handle updated data from the coordinator.""" - self._attr_is_on = self._device.status + self._async_update_attrs() super()._handle_coordinator_update() + @callback + def _async_update_attrs(self) -> None: + """Handle updating _attr values.""" + self._attr_is_on = self._device.status + async def async_turn_off(self, **kwargs: Any) -> None: """Instruct the device to turn off.""" await self._device.turn_off() diff --git a/homeassistant/components/wiz/light.py b/homeassistant/components/wiz/light.py index ee586c8939d..bc5ac078ec7 100644 --- a/homeassistant/components/wiz/light.py +++ b/homeassistant/components/wiz/light.py @@ -5,7 +5,7 @@ import logging from typing import Any from pywizlight import PilotBuilder -from pywizlight.bulblibrary import BulbClass, BulbType +from pywizlight.bulblibrary import BulbClass, BulbType, Features from pywizlight.rgbcw import convertHSfromRGBCW from pywizlight.scenes import get_id_from_scene_name @@ -34,34 +34,6 @@ from .models import WizData _LOGGER = logging.getLogger(__name__) -DEFAULT_COLOR_MODES = {COLOR_MODE_HS, COLOR_MODE_COLOR_TEMP} -DEFAULT_MIN_MIREDS = 153 -DEFAULT_MAX_MIREDS = 454 - - -def get_supported_color_modes(bulb_type: BulbType) -> set[str]: - """Flag supported features.""" - color_modes = set() - features = bulb_type.features - if features.color: - color_modes.add(COLOR_MODE_HS) - if features.color_tmp: - color_modes.add(COLOR_MODE_COLOR_TEMP) - if not color_modes and features.brightness: - color_modes.add(COLOR_MODE_BRIGHTNESS) - return color_modes - - -def get_min_max_mireds(bulb_type: BulbType) -> tuple[int, int]: - """Return the coldest and warmest color_temp that this light supports.""" - # DW bulbs have no kelvin - if bulb_type.bulb_type == BulbClass.DW: - return 0, 0 - # If bulbtype is TW or RGB then return the kelvin value - return color_temperature_kelvin_to_mired( - bulb_type.kelvin_range.max - ), color_temperature_kelvin_to_mired(bulb_type.kelvin_range.min) - async def async_setup_entry( hass: HomeAssistant, @@ -81,20 +53,35 @@ class WizBulbEntity(WizToggleEntity, LightEntity): """Initialize an WiZLight.""" super().__init__(wiz_data, name) bulb_type: BulbType = self._device.bulbtype + features: Features = bulb_type.features + color_modes = set() + if features.color: + color_modes.add(COLOR_MODE_HS) + if features.color_tmp: + color_modes.add(COLOR_MODE_COLOR_TEMP) + if not color_modes and features.brightness: + color_modes.add(COLOR_MODE_BRIGHTNESS) + self._attr_supported_color_modes = color_modes self._attr_effect_list = wiz_data.scenes - self._attr_min_mireds, self._attr_max_mireds = get_min_max_mireds(bulb_type) - self._attr_supported_color_modes = get_supported_color_modes(bulb_type) + if bulb_type.bulb_type != BulbClass.DW: + self._attr_min_mireds = color_temperature_kelvin_to_mired( + bulb_type.kelvin_range.max + ) + self._attr_max_mireds = color_temperature_kelvin_to_mired( + bulb_type.kelvin_range.min + ) if bulb_type.features.effect: self._attr_supported_features = SUPPORT_EFFECT + self._async_update_attrs() @callback - def _handle_coordinator_update(self) -> None: - """Handle updated data from the coordinator.""" + def _async_update_attrs(self) -> None: + """Handle updating _attr values.""" state = self._device.state - if (brightness := state.get_brightness()) is not None: - self._attr_brightness = max(0, min(255, brightness)) color_modes = self.supported_color_modes assert color_modes is not None + if (brightness := state.get_brightness()) is not None: + self._attr_brightness = max(0, min(255, brightness)) if COLOR_MODE_COLOR_TEMP in color_modes and state.get_colortemp() is not None: self._attr_color_mode = COLOR_MODE_COLOR_TEMP if color_temp := state.get_colortemp(): @@ -110,7 +97,7 @@ class WizBulbEntity(WizToggleEntity, LightEntity): else: self._attr_color_mode = COLOR_MODE_BRIGHTNESS self._attr_effect = state.get_scene() - super()._handle_coordinator_update() + super()._async_update_attrs() @callback def _async_pilot_builder(self, **kwargs: Any) -> PilotBuilder: diff --git a/homeassistant/components/wiz/strings.json b/homeassistant/components/wiz/strings.json index 2195bb09a03..288fd76acc4 100644 --- a/homeassistant/components/wiz/strings.json +++ b/homeassistant/components/wiz/strings.json @@ -6,10 +6,15 @@ "data": { "host": "[%key:common::config_flow::data::host%]" }, - "description": "Enter the IP address of the device." + "description": "If you leave the host empty, discovery will be used to find devices." }, "discovery_confirm": { "description": "Do you want to setup {name} ({host})?" + }, + "pick_device": { + "data": { + "device": "Device" + } } }, "error": { diff --git a/homeassistant/components/wiz/switch.py b/homeassistant/components/wiz/switch.py index eae7e3d47a8..e6e34c73c3b 100644 --- a/homeassistant/components/wiz/switch.py +++ b/homeassistant/components/wiz/switch.py @@ -29,6 +29,11 @@ async def async_setup_entry( class WizSocketEntity(WizToggleEntity, SwitchEntity): """Representation of a WiZ socket.""" + def __init__(self, wiz_data: WizData, name: str) -> None: + """Initialize a WiZ socket.""" + super().__init__(wiz_data, name) + self._async_update_attrs() + async def async_turn_on(self, **kwargs: Any) -> None: """Instruct the socket to turn on.""" await self._device.turn_on(PilotBuilder()) diff --git a/homeassistant/components/wiz/translations/en.json b/homeassistant/components/wiz/translations/en.json index 192d1137c2f..c8ee3d6df2e 100644 --- a/homeassistant/components/wiz/translations/en.json +++ b/homeassistant/components/wiz/translations/en.json @@ -14,11 +14,16 @@ "discovery_confirm": { "description": "Do you want to setup {name} ({host})?" }, + "pick_device": { + "data": { + "device": "Device" + } + }, "user": { "data": { "host": "Host" }, - "description": "Enter the IP address of the device." + "description": "If you leave the host empty, discovery will be used to find devices." } } } diff --git a/tests/components/wiz/test_config_flow.py b/tests/components/wiz/test_config_flow.py index a043d4a654e..a52ca323830 100644 --- a/tests/components/wiz/test_config_flow.py +++ b/tests/components/wiz/test_config_flow.py @@ -4,13 +4,12 @@ from copy import deepcopy from unittest.mock import patch import pytest +from pywizlight.discovery import DiscoveredBulb +from pywizlight.exceptions import WizLightConnectionError, WizLightTimeOutError from homeassistant import config_entries from homeassistant.components import dhcp -from homeassistant.components.wiz.config_flow import ( - WizLightConnectionError, - WizLightTimeOutError, -) +from homeassistant.components.wiz.config_flow import CONF_DEVICE from homeassistant.components.wiz.const import DOMAIN from homeassistant.const import CONF_HOST from homeassistant.data_entry_flow import RESULT_TYPE_ABORT, RESULT_TYPE_FORM @@ -74,6 +73,18 @@ def _patch_wizlight(device=None, extended_white_range=None): return _patcher() +def _patch_discovery(): + @contextmanager + def _patcher(): + with patch( + "homeassistant.components.wiz.discovery.find_wizlights", + return_value=[DiscoveredBulb(FAKE_IP, FAKE_MAC)], + ): + yield + + return _patcher() + + async def test_form(hass): """Test we get the form.""" result = await hass.config_entries.flow.async_init( @@ -85,7 +96,9 @@ async def test_form(hass): with _patch_wizlight(), patch( "homeassistant.components.wiz.async_setup_entry", return_value=True, - ) as mock_setup_entry: + ) as mock_setup_entry, patch( + "homeassistant.components.wiz.async_setup", return_value=True + ) as mock_setup: result2 = await hass.config_entries.flow.async_configure( result["flow_id"], TEST_CONNECTION, @@ -97,6 +110,7 @@ async def test_form(hass): assert result2["data"] == { CONF_HOST: "1.1.1.1", } + assert len(mock_setup.mock_calls) == 1 assert len(mock_setup_entry.mock_calls) == 1 @@ -140,10 +154,7 @@ async def test_form_updates_unique_id(hass): result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_USER} ) - with _patch_wizlight(), patch( - "homeassistant.components.wiz.async_setup_entry", - return_value=True, - ): + with _patch_wizlight(): result2 = await hass.config_entries.flow.async_configure( result["flow_id"], TEST_CONNECTION, @@ -226,7 +237,9 @@ async def test_discovered_by_dhcp_or_integration_discovery( with patch( "homeassistant.components.wiz.async_setup_entry", return_value=True, - ) as mock_setup_entry: + ) as mock_setup_entry, patch( + "homeassistant.components.wiz.async_setup", return_value=True + ) as mock_setup: result2 = await hass.config_entries.flow.async_configure( result["flow_id"], {}, @@ -238,6 +251,7 @@ async def test_discovered_by_dhcp_or_integration_discovery( assert result2["data"] == { CONF_HOST: "1.1.1.1", } + assert len(mock_setup.mock_calls) == 1 assert len(mock_setup_entry.mock_calls) == 1 @@ -268,3 +282,104 @@ async def test_discovered_by_dhcp_or_integration_discovery_updates_host( assert result["type"] == RESULT_TYPE_ABORT assert result["reason"] == "already_configured" assert entry.data[CONF_HOST] == FAKE_IP + + +async def test_setup_via_discovery(hass): + """Test setting up via discovery.""" + result = await hass.config_entries.flow.async_init( + DOMAIN, context={"source": config_entries.SOURCE_USER} + ) + await hass.async_block_till_done() + assert result["type"] == "form" + assert result["step_id"] == "user" + assert not result["errors"] + + with _patch_discovery(): + result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + await hass.async_block_till_done() + + assert result2["type"] == "form" + assert result2["step_id"] == "pick_device" + assert not result2["errors"] + + # test we can try again + result = await hass.config_entries.flow.async_init( + DOMAIN, context={"source": config_entries.SOURCE_USER} + ) + assert result["type"] == "form" + assert result["step_id"] == "user" + assert not result["errors"] + + with _patch_discovery(): + result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + await hass.async_block_till_done() + + assert result2["type"] == "form" + assert result2["step_id"] == "pick_device" + assert not result2["errors"] + + with _patch_wizlight(), patch( + "homeassistant.components.wiz.async_setup", return_value=True + ) as mock_setup, patch( + "homeassistant.components.wiz.async_setup_entry", return_value=True + ) as mock_setup_entry: + result3 = await hass.config_entries.flow.async_configure( + result["flow_id"], + {CONF_DEVICE: FAKE_MAC}, + ) + await hass.async_block_till_done() + + assert result3["type"] == "create_entry" + assert result3["title"] == "WiZ Dimmable White ABCABC" + assert result3["data"] == { + CONF_HOST: "1.1.1.1", + } + assert len(mock_setup.mock_calls) == 1 + assert len(mock_setup_entry.mock_calls) == 1 + + # ignore configured devices + result = await hass.config_entries.flow.async_init( + DOMAIN, context={"source": config_entries.SOURCE_USER} + ) + assert result["type"] == "form" + assert result["step_id"] == "user" + assert not result["errors"] + + with _patch_discovery(): + result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + await hass.async_block_till_done() + + assert result2["type"] == "abort" + assert result2["reason"] == "no_devices_found" + + +async def test_setup_via_discovery_cannot_connect(hass): + """Test setting up via discovery and we fail to connect to the discovered device.""" + result = await hass.config_entries.flow.async_init( + DOMAIN, context={"source": config_entries.SOURCE_USER} + ) + await hass.async_block_till_done() + assert result["type"] == "form" + assert result["step_id"] == "user" + assert not result["errors"] + + with _patch_discovery(): + result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + await hass.async_block_till_done() + + assert result2["type"] == "form" + assert result2["step_id"] == "pick_device" + assert not result2["errors"] + + with patch( + "homeassistant.components.wiz.wizlight.getBulbConfig", + side_effect=WizLightTimeOutError, + ), _patch_discovery(): + result3 = await hass.config_entries.flow.async_configure( + result["flow_id"], + {CONF_DEVICE: FAKE_MAC}, + ) + await hass.async_block_till_done() + + assert result3["type"] == "abort" + assert result3["reason"] == "cannot_connect"