Add discovery support to WiZ Part 1 (#65752)

.coveragerc

@@ -1316,6 +1316,7 @@ omit =
     homeassistant/components/wirelesstag/*
     homeassistant/components/wiz/__init__.py
     homeassistant/components/wiz/const.py
+    homeassistant/components/wiz/discovery.py
     homeassistant/components/wiz/light.py
     homeassistant/components/wolflink/__init__.py
     homeassistant/components/wolflink/sensor.py

homeassistant/components/wiz/__init__.py

@@ -1,17 +1,23 @@
 """WiZ Platform integration."""
+import asyncio
 from datetime import timedelta
 import logging
+from typing import Any
 
 from pywizlight import wizlight
+from pywizlight.exceptions import WizLightNotKnownBulb
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.debounce import Debouncer
+from homeassistant.helpers.event import async_track_time_interval
+from homeassistant.helpers.typing import ConfigType
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import DOMAIN, WIZ_EXCEPTIONS
+from .const import DISCOVER_SCAN_TIMEOUT, DISCOVERY_INTERVAL, DOMAIN, WIZ_EXCEPTIONS
+from .discovery import async_discover_devices, async_trigger_discovery
 from .models import WizData
 
 _LOGGER = logging.getLogger(__name__)
@@ -21,6 +27,19 @@ PLATFORMS = [Platform.LIGHT]
 REQUEST_REFRESH_DELAY = 0.35
 
 
+async def async_setup(hass: HomeAssistant, hass_config: ConfigType) -> bool:
+    """Set up the wiz integration."""
+
+    async def _async_discovery(*_: Any) -> None:
+        async_trigger_discovery(
+            hass, await async_discover_devices(hass, DISCOVER_SCAN_TIMEOUT)
+        )
+
+    asyncio.create_task(_async_discovery())
+    async_track_time_interval(hass, _async_discovery, DISCOVERY_INTERVAL)
+    return True
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up the wiz integration from a config entry."""
     ip_address = entry.data[CONF_HOST]
@@ -34,6 +53,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         # This is likely because way the library
         # processes responses and can be cleaned up
         # in the future.
+    except WizLightNotKnownBulb:
+        # This is only thrown on IndexError when the
+        # bulb responds with invalid data? It may
+        # not actually be possible anymore
+        _LOGGER.warning("The WiZ bulb type could not be determined for %s", ip_address)
+        return False
     except (ValueError, *WIZ_EXCEPTIONS) as err:
         raise ConfigEntryNotReady from err
 

homeassistant/components/wiz/config_flow.py

@@ -5,15 +5,17 @@ import logging
 from typing import Any
 
 from pywizlight import wizlight
+from pywizlight.discovery import DiscoveredBulb
 from pywizlight.exceptions import WizLightConnectionError, WizLightTimeOutError
 import voluptuous as vol
 
 from homeassistant import config_entries
+from homeassistant.components import dhcp
 from homeassistant.const import CONF_HOST
 from homeassistant.data_entry_flow import FlowResult
 
-from .const import DEFAULT_NAME, DOMAIN
-from .utils import _short_mac
+from .const import DOMAIN, WIZ_EXCEPTIONS
+from .utils import name_from_bulb_type_and_mac
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -23,6 +25,66 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
     VERSION = 1
 
+    def __init__(self) -> None:
+        """Initialize the config flow."""
+        self._discovered_device: DiscoveredBulb | None = None
+        self._name: str | None = None
+
+    async def async_step_dhcp(self, discovery_info: dhcp.DhcpServiceInfo) -> FlowResult:
+        """Handle discovery via dhcp."""
+        self._discovered_device = DiscoveredBulb(
+            discovery_info.ip, discovery_info.macaddress
+        )
+        return await self._async_handle_discovery()
+
+    async def async_step_integration_discovery(
+        self, discovery_info: dict[str, str]
+    ) -> FlowResult:
+        """Handle integration discovery."""
+        self._discovered_device = DiscoveredBulb(
+            discovery_info["ip_address"], discovery_info["mac_address"]
+        )
+        return await self._async_handle_discovery()
+
+    async def _async_handle_discovery(self) -> FlowResult:
+        """Handle any discovery."""
+        device = self._discovered_device
+        assert device is not None
+        _LOGGER.debug("Discovered device: %s", device)
+        ip_address = device.ip_address
+        mac = device.mac_address
+        await self.async_set_unique_id(mac)
+        self._abort_if_unique_id_configured(updates={CONF_HOST: ip_address})
+        bulb = wizlight(ip_address)
+        try:
+            bulbtype = await bulb.get_bulbtype()
+        except WIZ_EXCEPTIONS:
+            return self.async_abort(reason="cannot_connect")
+        self._name = name_from_bulb_type_and_mac(bulbtype, mac)
+        return await self.async_step_discovery_confirm()
+
+    async def async_step_discovery_confirm(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Confirm discovery."""
+        assert self._discovered_device is not None
+        assert self._name is not None
+        ip_address = self._discovered_device.ip_address
+        if user_input is not None:
+            return self.async_create_entry(
+                title=self._name,
+                data={CONF_HOST: ip_address},
+            )
+
+        self._set_confirm_only()
+        placeholders = {"name": self._name, "host": ip_address}
+        self.context["title_placeholders"] = placeholders
+        return self.async_show_form(
+            step_id="discovery_confirm",
+            description_placeholders=placeholders,
+            data_schema=vol.Schema({}),
+        )
+
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
     ) -> FlowResult:
@@ -43,12 +105,11 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 _LOGGER.exception("Unexpected exception")
                 errors["base"] = "unknown"
             else:
-                await self.async_set_unique_id(mac)
+                await self.async_set_unique_id(mac, raise_on_progress=False)
                 self._abort_if_unique_id_configured(
                     updates={CONF_HOST: user_input[CONF_HOST]}
                 )
-                bulb_type = bulbtype.bulb_type.value if bulbtype else "Unknown"
-                name = f"{DEFAULT_NAME} {bulb_type} {_short_mac(mac)}"
+                name = name_from_bulb_type_and_mac(bulbtype, mac)
                 return self.async_create_entry(
                     title=name,
                     data=user_input,

homeassistant/components/wiz/const.py

@@ -1,9 +1,16 @@
 """Constants for the WiZ Platform integration."""
+from datetime import timedelta
+
 from pywizlight.exceptions import WizLightConnectionError, WizLightTimeOutError
 
 DOMAIN = "wiz"
 DEFAULT_NAME = "WiZ"
 
+DISCOVER_SCAN_TIMEOUT = 10
+DISCOVERY_INTERVAL = timedelta(minutes=15)
+
+SOCKET_DEVICE_STR = "_SOCKET_"
+
 WIZ_EXCEPTIONS = (
     OSError,
     WizLightTimeOutError,

homeassistant/components/wiz/discovery.py

@@ -0,0 +1,61 @@
+"""The wiz integration discovery."""
+from __future__ import annotations
+
+import asyncio
+from dataclasses import asdict
+import logging
+
+from pywizlight.discovery import DiscoveredBulb, find_wizlights
+
+from homeassistant import config_entries
+from homeassistant.components import network
+from homeassistant.core import HomeAssistant, callback
+
+from .const import DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_discover_devices(
+    hass: HomeAssistant, timeout: int, address: str | None = None
+) -> list[DiscoveredBulb]:
+    """Discover wiz devices."""
+    if address:
+        targets = [address]
+    else:
+        targets = [
+            str(address)
+            for address in await network.async_get_ipv4_broadcast_addresses(hass)
+        ]
+
+    combined_discoveries: dict[str, DiscoveredBulb] = {}
+    for idx, discovered in enumerate(
+        await asyncio.gather(
+            *[find_wizlights(timeout, address) for address in targets],
+            return_exceptions=True,
+        )
+    ):
+        if isinstance(discovered, Exception):
+            _LOGGER.debug("Scanning %s failed with error: %s", targets[idx], discovered)
+            continue
+        for device in discovered:
+            assert isinstance(device, DiscoveredBulb)
+            combined_discoveries[device.ip_address] = device
+
+    return list(combined_discoveries.values())
+
+
+@callback
+def async_trigger_discovery(
+    hass: HomeAssistant,
+    discovered_devices: list[DiscoveredBulb],
+) -> None:
+    """Trigger config flows for discovered devices."""
+    for device in discovered_devices:
+        hass.async_create_task(
+            hass.config_entries.flow.async_init(
+                DOMAIN,
+                context={"source": config_entries.SOURCE_INTEGRATION_DISCOVERY},
+                data=asdict(device),
+            )
+        )

homeassistant/components/wiz/light.py

@@ -1,13 +1,11 @@
 """WiZ integration."""
 from __future__ import annotations
 
-import contextlib
 import logging
 from typing import Any
 
 from pywizlight import PilotBuilder
 from pywizlight.bulblibrary import BulbClass, BulbType
-from pywizlight.exceptions import WizLightNotKnownBulb
 from pywizlight.rgbcw import convertHSfromRGBCW
 from pywizlight.scenes import get_id_from_scene_name
 
@@ -43,29 +41,20 @@ DEFAULT_MAX_MIREDS = 454
 
 def get_supported_color_modes(bulb_type: BulbType) -> set[str]:
     """Flag supported features."""
-    if not bulb_type:
-        # fallback
-        return DEFAULT_COLOR_MODES
     color_modes = set()
-    try:
-        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
-    except WizLightNotKnownBulb:
-        _LOGGER.warning("Bulb is not present in the library. Fallback to full feature")
-        return DEFAULT_COLOR_MODES
+    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 supports_effects(bulb_type: BulbType) -> bool:
     """Check if a bulb supports effects."""
-    with contextlib.suppress(WizLightNotKnownBulb):
-        return bool(bulb_type.features.effect)
-    return True  # default is true
+    return bool(bulb_type.features.effect)
 
 
 def get_min_max_mireds(bulb_type: BulbType) -> tuple[int, int]:
@@ -76,13 +65,9 @@ def get_min_max_mireds(bulb_type: BulbType) -> tuple[int, int]:
     if bulb_type.bulb_type == BulbClass.DW:
         return 0, 0
     # If bulbtype is TW or RGB then return the kelvin value
-    try:
-        return color_utils.color_temperature_kelvin_to_mired(
-            bulb_type.kelvin_range.max
-        ), color_utils.color_temperature_kelvin_to_mired(bulb_type.kelvin_range.min)
-    except WizLightNotKnownBulb:
-        _LOGGER.debug("Kelvin is not present in the library. Fallback to 6500")
-        return DEFAULT_MIN_MIREDS, DEFAULT_MAX_MIREDS
+    return color_utils.color_temperature_kelvin_to_mired(
+        bulb_type.kelvin_range.max
+    ), color_utils.color_temperature_kelvin_to_mired(bulb_type.kelvin_range.min)
 
 
 async def async_setup_entry(

homeassistant/components/wiz/manifest.json

@@ -2,8 +2,13 @@
   "domain": "wiz",
   "name": "WiZ",
   "config_flow": true,
+  "dhcp": [
+    {"macaddress":"A8BB50*"},
+    {"hostname":"wiz_[0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f]"}
+  ],
+  "dependencies": ["network"],
   "documentation": "https://www.home-assistant.io/integrations/wiz",
   "requirements": ["pywizlight==0.4.16"],
   "iot_class": "local_polling",
   "codeowners": ["@sbidy"]
-}
+}

homeassistant/components/wiz/strings.json

@@ -1,11 +1,15 @@
 {
   "config": {
+    "flow_title": "{name} ({host})",
     "step": {
       "user": {
         "data": {
           "host": "[%key:common::config_flow::data::host%]"
         },
         "description": "Enter the IP address of the device."
+      },
+      "discovery_confirm": {
+        "description": "Do you want to setup {name} ({host})?"
       }
     },
     "error": {

homeassistant/components/wiz/translations/en.json

@@ -10,7 +10,11 @@
             "no_wiz_light": "The bulb can not be connected via WiZ Platform integration.",
             "unknown": "Unexpected error"
         },
+        "flow_title": "{name} ({host})",
         "step": {
+            "discovery_confirm": {
+                "description": "Do you want to setup {name} ({host})?"
+            },
             "user": {
                 "data": {
                     "host": "Host"

homeassistant/components/wiz/utils.py

@@ -1,7 +1,20 @@
 """WiZ utils."""
 from __future__ import annotations
 
+from pywizlight import BulbType
+
+from .const import DEFAULT_NAME, SOCKET_DEVICE_STR
+
 
 def _short_mac(mac: str) -> str:
     """Get the short mac address from the full mac."""
     return mac.replace(":", "").upper()[-6:]
+
+
+def name_from_bulb_type_and_mac(bulb_type: BulbType, mac: str) -> str:
+    """Generate a name from bulb_type and mac."""
+    if SOCKET_DEVICE_STR in bulb_type.name:
+        description = "Socket"
+    else:
+        description = bulb_type.bulb_type.value
+    return f"{DEFAULT_NAME} {description} {_short_mac(mac)}"

homeassistant/generated/dhcp.py

@@ -608,6 +608,14 @@ DHCP = [
         "domain": "vicare",
         "macaddress": "B87424*"
     },
+    {
+        "domain": "wiz",
+        "macaddress": "A8BB50*"
+    },
+    {
+        "domain": "wiz",
+        "hostname": "wiz_[0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f]"
+    },
     {
         "domain": "yeelight",
         "hostname": "yeelink-*"

tests/components/wiz/test_config_flow.py

@@ -1,19 +1,23 @@
 """Test the WiZ Platform config flow."""
 from contextlib import contextmanager
+from copy import deepcopy
 from unittest.mock import patch
 
 import pytest
 
 from homeassistant import config_entries
+from homeassistant.components import dhcp
 from homeassistant.components.wiz.config_flow import (
     WizLightConnectionError,
     WizLightTimeOutError,
 )
 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
 
 from tests.common import MockConfigEntry
 
+FAKE_IP = "1.1.1.1"
 FAKE_MAC = "ABCABCABCABC"
 FAKE_BULB_CONFIG = {
     "method": "getSystemConfig",
@@ -31,21 +35,36 @@ FAKE_BULB_CONFIG = {
         "ping": 0,
     },
 }
+FAKE_SOCKET_CONFIG = deepcopy(FAKE_BULB_CONFIG)
+FAKE_SOCKET_CONFIG["result"]["moduleName"] = "ESP10_SOCKET_06"
 FAKE_EXTENDED_WHITE_RANGE = [2200, 2700, 6500, 6500]
 TEST_SYSTEM_INFO = {"id": FAKE_MAC, "name": "Test Bulb"}
 TEST_CONNECTION = {CONF_HOST: "1.1.1.1"}
 TEST_NO_IP = {CONF_HOST: "this is no IP input"}
 
 
-def _patch_wizlight():
+DHCP_DISCOVERY = dhcp.DhcpServiceInfo(
+    hostname="wiz_abcabc",
+    ip=FAKE_IP,
+    macaddress=FAKE_MAC,
+)
+
+
+INTEGRATION_DISCOVERY = {
+    "ip_address": FAKE_IP,
+    "mac_address": FAKE_MAC,
+}
+
+
+def _patch_wizlight(device=None, extended_white_range=None):
     @contextmanager
     def _patcher():
         with patch(
             "homeassistant.components.wiz.wizlight.getBulbConfig",
-            return_value=FAKE_BULB_CONFIG,
+            return_value=device or FAKE_BULB_CONFIG,
         ), patch(
             "homeassistant.components.wiz.wizlight.getExtendedWhiteRange",
-            return_value=FAKE_EXTENDED_WHITE_RANGE,
+            return_value=extended_white_range or FAKE_EXTENDED_WHITE_RANGE,
         ), patch(
             "homeassistant.components.wiz.wizlight.getMac",
             return_value=FAKE_MAC,
@@ -114,11 +133,8 @@ async def test_form_updates_unique_id(hass):
     entry = MockConfigEntry(
         domain=DOMAIN,
         unique_id=TEST_SYSTEM_INFO["id"],
-        data={
-            CONF_HOST: "dummy",
-        },
+        data={CONF_HOST: "dummy"},
     )
-
     entry.add_to_hass(hass)
 
     result = await hass.config_entries.flow.async_init(
@@ -136,3 +152,118 @@ async def test_form_updates_unique_id(hass):
 
     assert result2["type"] == "abort"
     assert result2["reason"] == "already_configured"
+
+
+@pytest.mark.parametrize(
+    "source, data",
+    [
+        (config_entries.SOURCE_DHCP, DHCP_DISCOVERY),
+        (config_entries.SOURCE_INTEGRATION_DISCOVERY, INTEGRATION_DISCOVERY),
+    ],
+)
+async def test_discovered_by_dhcp_connection_fails(hass, source, data):
+    """Test we abort on connection failure."""
+    with patch(
+        "homeassistant.components.wiz.wizlight.getBulbConfig",
+        side_effect=WizLightTimeOutError,
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": source}, data=data
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == RESULT_TYPE_ABORT
+    assert result["reason"] == "cannot_connect"
+
+
+@pytest.mark.parametrize(
+    "source, data, device, extended_white_range, name",
+    [
+        (
+            config_entries.SOURCE_DHCP,
+            DHCP_DISCOVERY,
+            FAKE_BULB_CONFIG,
+            FAKE_EXTENDED_WHITE_RANGE,
+            "WiZ Dimmable White ABCABC",
+        ),
+        (
+            config_entries.SOURCE_INTEGRATION_DISCOVERY,
+            INTEGRATION_DISCOVERY,
+            FAKE_BULB_CONFIG,
+            FAKE_EXTENDED_WHITE_RANGE,
+            "WiZ Dimmable White ABCABC",
+        ),
+        (
+            config_entries.SOURCE_DHCP,
+            DHCP_DISCOVERY,
+            FAKE_SOCKET_CONFIG,
+            None,
+            "WiZ Socket ABCABC",
+        ),
+        (
+            config_entries.SOURCE_INTEGRATION_DISCOVERY,
+            INTEGRATION_DISCOVERY,
+            FAKE_SOCKET_CONFIG,
+            None,
+            "WiZ Socket ABCABC",
+        ),
+    ],
+)
+async def test_discovered_by_dhcp_or_integration_discovery(
+    hass, source, data, device, extended_white_range, name
+):
+    """Test we can configure when discovered from dhcp or discovery."""
+    with _patch_wizlight(device=device, extended_white_range=extended_white_range):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": source}, data=data
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["step_id"] == "discovery_confirm"
+
+    with patch(
+        "homeassistant.components.wiz.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry:
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {},
+        )
+        await hass.async_block_till_done()
+
+    assert result2["type"] == "create_entry"
+    assert result2["title"] == name
+    assert result2["data"] == {
+        CONF_HOST: "1.1.1.1",
+    }
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
+@pytest.mark.parametrize(
+    "source, data",
+    [
+        (config_entries.SOURCE_DHCP, DHCP_DISCOVERY),
+        (config_entries.SOURCE_INTEGRATION_DISCOVERY, INTEGRATION_DISCOVERY),
+    ],
+)
+async def test_discovered_by_dhcp_or_integration_discovery_updates_host(
+    hass, source, data
+):
+    """Test dhcp or discovery updates existing host."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        unique_id=TEST_SYSTEM_INFO["id"],
+        data={CONF_HOST: "dummy"},
+    )
+    entry.add_to_hass(hass)
+
+    with _patch_wizlight():
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": source}, data=data
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == RESULT_TYPE_ABORT
+    assert result["reason"] == "already_configured"
+    assert entry.data[CONF_HOST] == FAKE_IP