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

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

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()

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:

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": {

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())

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."
             }
         }
     }

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"