From 9229d14962030f0a5fd11941c2533aa8adb27bac Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Wed, 22 Jun 2022 20:17:28 +0200
Subject: [PATCH] Automatically onboard Wiz (#73851)

---
 homeassistant/components/wiz/config_flow.py | 13 ++++-----
 tests/components/wiz/test_config_flow.py    | 31 +++++++++++++++++++++
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/homeassistant/components/wiz/config_flow.py b/homeassistant/components/wiz/config_flow.py
index b2173ccda97..b1bce3eda0d 100644
--- a/homeassistant/components/wiz/config_flow.py
+++ b/homeassistant/components/wiz/config_flow.py
@@ -9,7 +9,7 @@ from pywizlight.discovery import DiscoveredBulb
 from pywizlight.exceptions import WizLightConnectionError, WizLightTimeOutError
 import voluptuous as vol
 
-from homeassistant.components import dhcp
+from homeassistant.components import dhcp, onboarding
 from homeassistant.config_entries import ConfigFlow
 from homeassistant.const import CONF_HOST
 from homeassistant.data_entry_flow import AbortFlow, FlowResult
@@ -29,11 +29,12 @@ class WizConfigFlow(ConfigFlow, domain=DOMAIN):
 
     VERSION = 1
 
+    _discovered_device: DiscoveredBulb
+    _name: str
+
     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:
         """Handle discovery via dhcp."""
@@ -54,7 +55,6 @@ class WizConfigFlow(ConfigFlow, domain=DOMAIN):
     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
@@ -66,7 +66,6 @@ class WizConfigFlow(ConfigFlow, domain=DOMAIN):
     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()
@@ -84,10 +83,8 @@ class WizConfigFlow(ConfigFlow, domain=DOMAIN):
         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:
+        if user_input is not None or not onboarding.async_is_onboarded(self.hass):
             # Make sure the device is still there and
             # update the name if the firmware has auto
             # updated since discovery
diff --git a/tests/components/wiz/test_config_flow.py b/tests/components/wiz/test_config_flow.py
index 58b46bbea9d..f37f2ba21a0 100644
--- a/tests/components/wiz/test_config_flow.py
+++ b/tests/components/wiz/test_config_flow.py
@@ -506,3 +506,34 @@ async def test_discovery_with_firmware_update(hass):
     }
     assert len(mock_setup.mock_calls) == 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_during_onboarding(hass, source, data):
+    """Test dhcp or discovery during onboarding creates the config entry."""
+    with _patch_wizlight(), 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(
+        "homeassistant.components.onboarding.async_is_onboarded", return_value=False
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": source}, data=data
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == "create_entry"
+    assert result["title"] == "WiZ Dimmable White ABCABC"
+    assert result["data"] == {
+        CONF_HOST: "1.1.1.1",
+    }
+    assert len(mock_setup.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 1