Retry right away on discovery for WiZ (#72659)

homeassistant/components/wiz/config_flow.py

@@ -9,8 +9,8 @@ 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.config_entries import ConfigEntryState, ConfigFlow
 from homeassistant.const import CONF_HOST
 from homeassistant.data_entry_flow import AbortFlow, FlowResult
 from homeassistant.util.network import is_ip_address
@@ -24,7 +24,7 @@ _LOGGER = logging.getLogger(__name__)
 CONF_DEVICE = "device"
 
 
-class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
+class WizConfigFlow(ConfigFlow, domain=DOMAIN):
     """Handle a config flow for WiZ."""
 
     VERSION = 1
@@ -58,7 +58,15 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         _LOGGER.debug("Discovered device: %s", device)
         ip_address = device.ip_address
         mac = device.mac_address
-        await self.async_set_unique_id(mac)
+        if current_entry := await self.async_set_unique_id(mac):
+            if (
+                current_entry.state is ConfigEntryState.SETUP_RETRY
+                and current_entry.data[CONF_HOST] == ip_address
+            ):
+                self.hass.async_create_task(
+                    self.hass.config_entries.async_reload(current_entry.entry_id)
+                )
+                return self.async_abort(reason="already_configured")
         self._abort_if_unique_id_configured(updates={CONF_HOST: ip_address})
         await self._async_connect_discovered_or_abort()
         return await self.async_step_discovery_confirm()

tests/components/wiz/test_config_flow.py

@@ -1,5 +1,5 @@
 """Test the WiZ Platform config flow."""
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
 import pytest
 from pywizlight.exceptions import WizLightConnectionError, WizLightTimeOutError
@@ -21,8 +21,10 @@ from . import (
     FAKE_SOCKET,
     TEST_CONNECTION,
     TEST_SYSTEM_INFO,
+    _mocked_wizlight,
     _patch_discovery,
     _patch_wizlight,
+    async_setup_integration,
 )
 
 from tests.common import MockConfigEntry
@@ -309,6 +311,35 @@ async def test_discovered_by_dhcp_or_integration_discovery_updates_host(
     assert entry.data[CONF_HOST] == FAKE_IP
 
 
+@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_avoid_waiting_for_retry(
+    hass, source, data
+):
+    """Test dhcp or discovery kicks off setup when in retry."""
+    bulb = _mocked_wizlight(None, None, FAKE_SOCKET)
+    bulb.getMac = AsyncMock(side_effect=OSError)
+    _, entry = await async_setup_integration(hass, wizlight=bulb)
+    assert entry.data[CONF_HOST] == FAKE_IP
+    assert entry.state is config_entries.ConfigEntryState.SETUP_RETRY
+    bulb.getMac = AsyncMock(return_value=FAKE_MAC)
+
+    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.state is config_entries.ConfigEntryState.LOADED
+
+
 async def test_setup_via_discovery(hass):
     """Test setting up via discovery."""
     result = await hass.config_entries.flow.async_init(

tests/components/wiz/test_init.py

@@ -1,13 +1,16 @@
 """Tests for wiz integration."""
 import datetime
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, patch
 
 from homeassistant import config_entries
-from homeassistant.const import ATTR_FRIENDLY_NAME, EVENT_HOMEASSISTANT_STOP
+from homeassistant.components.wiz.const import DOMAIN
+from homeassistant.const import ATTR_FRIENDLY_NAME, CONF_HOST, EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
 from homeassistant.util.dt import utcnow
 
 from . import (
+    FAKE_IP,
     FAKE_MAC,
     FAKE_SOCKET,
     _mocked_wizlight,
@@ -16,7 +19,7 @@ from . import (
     async_setup_integration,
 )
 
-from tests.common import async_fire_time_changed
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 
 async def test_setup_retry(hass: HomeAssistant) -> None:
@@ -47,7 +50,17 @@ async def test_cleanup_on_failed_first_update(hass: HomeAssistant) -> None:
     """Test the socket is cleaned up on failed first update."""
     bulb = _mocked_wizlight(None, None, FAKE_SOCKET)
     bulb.updateState = AsyncMock(side_effect=OSError)
-    _, entry = await async_setup_integration(hass, wizlight=bulb)
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        unique_id=FAKE_MAC,
+        data={CONF_HOST: FAKE_IP},
+    )
+    entry.add_to_hass(hass)
+    with patch(
+        "homeassistant.components.wiz.discovery.find_wizlights", return_value=[]
+    ), _patch_wizlight(device=bulb):
+        await async_setup_component(hass, DOMAIN, {DOMAIN: {}})
+        await hass.async_block_till_done()
     assert entry.state is config_entries.ConfigEntryState.SETUP_RETRY
     bulb.async_close.assert_called_once()