From 1ead6d6762027f8706b58f080149f9688d104dd0 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Wed, 22 Jun 2022 21:19:34 +0200
Subject: [PATCH] Automatically onboard Yeelight (#73854)

---
 .../components/yeelight/config_flow.py        |  4 +-
 tests/components/yeelight/test_config_flow.py | 50 +++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/yeelight/config_flow.py b/homeassistant/components/yeelight/config_flow.py
index 5a457b7da95..b4afedd6c51 100644
--- a/homeassistant/components/yeelight/config_flow.py
+++ b/homeassistant/components/yeelight/config_flow.py
@@ -11,7 +11,7 @@ from yeelight.aio import AsyncBulb
 from yeelight.main import get_known_models
 
 from homeassistant import config_entries, exceptions
-from homeassistant.components import dhcp, ssdp, zeroconf
+from homeassistant.components import dhcp, onboarding, ssdp, zeroconf
 from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.const import CONF_DEVICE, CONF_HOST, CONF_ID, CONF_MODEL, CONF_NAME
 from homeassistant.core import callback
@@ -134,7 +134,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
     async def async_step_discovery_confirm(self, user_input=None):
         """Confirm discovery."""
-        if user_input is not None:
+        if user_input is not None or not onboarding.async_is_onboarded(self.hass):
             return self.async_create_entry(
                 title=async_format_model_id(self._discovered_model, self.unique_id),
                 data={
diff --git a/tests/components/yeelight/test_config_flow.py b/tests/components/yeelight/test_config_flow.py
index 1c19a5e7dfd..7d9acc670b5 100644
--- a/tests/components/yeelight/test_config_flow.py
+++ b/tests/components/yeelight/test_config_flow.py
@@ -809,3 +809,53 @@ async def test_discovery_adds_missing_ip_id_only(hass: HomeAssistant):
     assert result["type"] == RESULT_TYPE_ABORT
     assert result["reason"] == "already_configured"
     assert config_entry.data[CONF_HOST] == IP_ADDRESS
+
+
+@pytest.mark.parametrize(
+    "source, data",
+    [
+        (
+            config_entries.SOURCE_DHCP,
+            dhcp.DhcpServiceInfo(
+                ip=IP_ADDRESS, macaddress="aa:bb:cc:dd:ee:ff", hostname="mock_hostname"
+            ),
+        ),
+        (
+            config_entries.SOURCE_HOMEKIT,
+            zeroconf.ZeroconfServiceInfo(
+                host=IP_ADDRESS,
+                addresses=[IP_ADDRESS],
+                hostname="mock_hostname",
+                name="mock_name",
+                port=None,
+                properties={zeroconf.ATTR_PROPERTIES_ID: "aa:bb:cc:dd:ee:ff"},
+                type="mock_type",
+            ),
+        ),
+    ],
+)
+async def test_discovered_during_onboarding(hass, source, data):
+    """Test we create a config entry when discovered during onboarding."""
+    mocked_bulb = _mocked_bulb()
+    with _patch_discovery(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ), patch(f"{MODULE}.async_setup", return_value=True) as mock_async_setup, patch(
+        f"{MODULE}.async_setup_entry", return_value=True
+    ) as mock_async_setup_entry, patch(
+        "homeassistant.components.onboarding.async_is_onboarded", return_value=False
+    ) as mock_is_onboarded:
+
+        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["data"] == {
+        CONF_HOST: IP_ADDRESS,
+        CONF_ID: "0x000000000015243f",
+        CONF_MODEL: MODEL,
+    }
+    assert mock_async_setup.called
+    assert mock_async_setup_entry.called
+    assert mock_is_onboarded.called