Add discovery on network up to WiZ (#66144)

homeassistant/components/wiz/__init__.py

@@ -83,6 +83,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
 
     await bulb.start_push(lambda _: coordinator.async_set_updated_data(None))
+    bulb.set_discovery_callback(lambda bulb: async_trigger_discovery(hass, [bulb]))
     await coordinator.async_config_entry_first_refresh()
 
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = WizData(

homeassistant/components/wiz/config_flow.py

@@ -12,7 +12,7 @@ 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 homeassistant.data_entry_flow import AbortFlow, FlowResult
 from homeassistant.util.network import is_ip_address
 
 from .const import DEFAULT_NAME, DISCOVER_SCAN_TIMEOUT, DOMAIN, WIZ_EXCEPTIONS
@@ -60,13 +60,19 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         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)
+        await self._async_connect_discovered_or_abort()
+        return await self.async_step_discovery_confirm()
+
+    async def _async_connect_discovered_or_abort(self) -> None:
+        """Connect to the device and verify its responding."""
+        device = self._discovered_device
+        assert device is not None
+        bulb = wizlight(device.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()
+        except WIZ_EXCEPTIONS as ex:
+            raise AbortFlow("cannot_connect") from ex
+        self._name = name_from_bulb_type_and_mac(bulbtype, device.mac_address)
 
     async def async_step_discovery_confirm(
         self, user_input: dict[str, Any] | None = None
@@ -76,6 +82,10 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         assert self._name is not None
         ip_address = self._discovered_device.ip_address
         if user_input is not None:
+            # Make sure the device is still there and
+            # update the name if the firmware has auto
+            # updated since discovery
+            await self._async_connect_discovered_or_abort()
             return self.async_create_entry(
                 title=self._name,
                 data={CONF_HOST: ip_address},

tests/components/wiz/test_config_flow.py

@@ -267,7 +267,9 @@ async def test_discovered_by_dhcp_or_integration_discovery(
     assert result["type"] == RESULT_TYPE_FORM
     assert result["step_id"] == "discovery_confirm"
 
-    with patch(
+    with _patch_wizlight(
+        device=device, extended_white_range=extended_white_range, bulb_type=bulb_type
+    ), patch(
         "homeassistant.components.wiz.async_setup_entry",
         return_value=True,
     ) as mock_setup_entry, patch(
@@ -416,3 +418,49 @@ async def test_setup_via_discovery_cannot_connect(hass):
 
     assert result3["type"] == "abort"
     assert result3["reason"] == "cannot_connect"
+
+
+async def test_discovery_with_firmware_update(hass):
+    """Test we check the device again between first discovery and config entry creation."""
+    with _patch_wizlight(
+        device=FAKE_BULB_CONFIG,
+        extended_white_range=FAKE_EXTENDED_WHITE_RANGE,
+        bulb_type=FAKE_RGBW_BULB,
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_INTEGRATION_DISCOVERY},
+            data=INTEGRATION_DISCOVERY,
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["step_id"] == "discovery_confirm"
+
+    # In between discovery and when the user clicks to set it up the firmware
+    # updates and we now can see its really RGBWW not RGBW since the older
+    # firmwares did not tell us how many white channels exist
+
+    with patch(
+        "homeassistant.components.wiz.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry, patch(
+        "homeassistant.components.wiz.async_setup", return_value=True
+    ) as mock_setup, _patch_wizlight(
+        device=FAKE_BULB_CONFIG,
+        extended_white_range=FAKE_EXTENDED_WHITE_RANGE,
+        bulb_type=FAKE_RGBWW_BULB,
+    ):
+        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"] == "WiZ RGBWW Tunable ABCABC"
+    assert result2["data"] == {
+        CONF_HOST: "1.1.1.1",
+    }
+    assert len(mock_setup.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 1