From a8374cf423cd9b5eab2631bb0b335cdc33356a64 Mon Sep 17 00:00:00 2001
From: Robert Svensson <Kane610@users.noreply.github.com>
Date: Thu, 30 Jan 2020 23:06:43 +0100
Subject: [PATCH] UniFi - Try to discover local controller (#31326)

* Its working

* Use "unifi" as default host if a controller can be found

* Fix tests

* Make a fixture of patching the discovery function
---
 homeassistant/components/unifi/config_flow.py | 16 +++++++++++++++-
 tests/components/unifi/conftest.py            | 13 +++++++++++++
 tests/components/unifi/test_config_flow.py    | 10 +++++++++-
 3 files changed, 37 insertions(+), 2 deletions(-)
 create mode 100644 tests/components/unifi/conftest.py

diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py
index 52ecab08856..9dbacc7916d 100644
--- a/homeassistant/components/unifi/config_flow.py
+++ b/homeassistant/components/unifi/config_flow.py
@@ -1,4 +1,6 @@
 """Config flow for UniFi."""
+import socket
+
 import voluptuous as vol
 
 from homeassistant import config_entries
@@ -104,11 +106,15 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
                 )
                 return self.async_abort(reason="unknown")
 
+        host = ""
+        if await async_discover_unifi(self.hass):
+            host = "unifi"
+
         return self.async_show_form(
             step_id="user",
             data_schema=vol.Schema(
                 {
-                    vol.Required(CONF_HOST): str,
+                    vol.Required(CONF_HOST, default=host): str,
                     vol.Required(CONF_USERNAME): str,
                     vol.Required(CONF_PASSWORD): str,
                     vol.Optional(CONF_PORT, default=DEFAULT_PORT): int,
@@ -235,3 +241,11 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
     async def _update_options(self):
         """Update config entry options."""
         return self.async_create_entry(title="", data=self.options)
+
+
+async def async_discover_unifi(hass):
+    """Discover UniFi address."""
+    try:
+        return await hass.async_add_executor_job(socket.gethostbyname, "unifi")
+    except socket.gaierror:
+        return None
diff --git a/tests/components/unifi/conftest.py b/tests/components/unifi/conftest.py
new file mode 100644
index 00000000000..189b80c1932
--- /dev/null
+++ b/tests/components/unifi/conftest.py
@@ -0,0 +1,13 @@
+"""Fixtures for UniFi methods."""
+from asynctest import patch
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def mock_discovery():
+    """No real network traffic allowed."""
+    with patch(
+        "homeassistant.components.unifi.config_flow.async_discover_unifi",
+        return_value=None,
+    ) as mock:
+        yield mock
diff --git a/tests/components/unifi/test_config_flow.py b/tests/components/unifi/test_config_flow.py
index cc8896d55ce..c2c89d2b9c0 100644
--- a/tests/components/unifi/test_config_flow.py
+++ b/tests/components/unifi/test_config_flow.py
@@ -16,14 +16,22 @@ from homeassistant.const import (
 from tests.common import MockConfigEntry
 
 
-async def test_flow_works(hass, aioclient_mock):
+async def test_flow_works(hass, aioclient_mock, mock_discovery):
     """Test config flow."""
+    mock_discovery.return_value = "1"
     result = await hass.config_entries.flow.async_init(
         config_flow.DOMAIN, context={"source": "user"}
     )
 
     assert result["type"] == "form"
     assert result["step_id"] == "user"
+    assert result["data_schema"]({CONF_USERNAME: "", CONF_PASSWORD: ""}) == {
+        CONF_HOST: "unifi",
+        CONF_USERNAME: "",
+        CONF_PASSWORD: "",
+        CONF_PORT: 8443,
+        CONF_VERIFY_SSL: False,
+    }
 
     aioclient_mock.post(
         "https://1.2.3.4:1234/api/login",