Add discovery to NUT integration (#36827)

* Add discovery to NUT integration

* implement async_step_zeroconf

* Update test to make sure unique id not set

* Remove host/port import when coming from discovery, add title placeholders

* fix mis-paste

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/nut/config_flow.py

@@ -36,14 +36,22 @@ SENSOR_DICT = {
     for sensor_id, sensor_spec in SENSOR_TYPES.items()
 }
 
-DATA_SCHEMA = vol.Schema(
+
+def _base_schema(discovery_info):
+    """Generate base schema."""
+    base_schema = {}
+    if not discovery_info:
+        base_schema.update(
             {
                 vol.Optional(CONF_HOST, default=DEFAULT_HOST): str,
                 vol.Optional(CONF_PORT, default=DEFAULT_PORT): int,
-        vol.Optional(CONF_USERNAME): str,
-        vol.Optional(CONF_PASSWORD): str,
             }
         )
+    base_schema.update(
+        {vol.Optional(CONF_USERNAME): str, vol.Optional(CONF_PASSWORD): str}
+    )
+
+    return vol.Schema(base_schema)
 
 
 def _resource_schema_base(available_resources, selected_resources):
@@ -75,7 +83,7 @@ def _ups_schema(ups_list):
 async def validate_input(hass: core.HomeAssistant, data):
     """Validate the user input allows us to connect.
 
-    Data has the keys from DATA_SCHEMA with values provided by the user.
+    Data has the keys from _base_schema with values provided by the user.
     """
 
     host = data[CONF_HOST]
@@ -113,9 +121,21 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         """Initialize the nut config flow."""
         self.nut_config = {}
         self.available_resources = {}
+        self.discovery_info = {}
         self.ups_list = None
         self.title = None
 
+    async def async_step_zeroconf(self, discovery_info):
+        """Prepare configuration for a discovered nut device."""
+        self.discovery_info = discovery_info
+        await self._async_handle_discovery_without_unique_id()
+        # pylint: disable=no-member # https://github.com/PyCQA/pylint/issues/3167
+        self.context["title_placeholders"] = {
+            CONF_PORT: discovery_info.get(CONF_PORT, DEFAULT_PORT),
+            CONF_HOST: discovery_info[CONF_HOST],
+        }
+        return await self.async_step_user()
+
     async def async_step_import(self, user_input=None):
         """Handle the import."""
         errors = {}
@@ -129,13 +149,20 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 return self.async_create_entry(title=title, data=user_input)
 
         return self.async_show_form(
-            step_id="user", data_schema=DATA_SCHEMA, errors=errors
+            step_id="user", data_schema=_base_schema({}), errors=errors
         )
 
     async def async_step_user(self, user_input=None):
         """Handle the user input."""
         errors = {}
         if user_input is not None:
+            if self.discovery_info:
+                user_input.update(
+                    {
+                        CONF_HOST: self.discovery_info[CONF_HOST],
+                        CONF_PORT: self.discovery_info.get(CONF_PORT, DEFAULT_PORT),
+                    }
+                )
             info, errors = await self._async_validate_or_error(user_input)
 
             if not errors:
@@ -150,7 +177,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 return await self.async_step_resources()
 
         return self.async_show_form(
-            step_id="user", data_schema=DATA_SCHEMA, errors=errors
+            step_id="user", data_schema=_base_schema(self.discovery_info), errors=errors
         )
 
     async def async_step_ups(self, user_input=None):

homeassistant/components/nut/manifest.json

@@ -4,5 +4,6 @@
   "documentation": "https://www.home-assistant.io/integrations/nut",
   "requirements": ["pynut2==2.1.2"],
   "codeowners": ["@bdraco"],
-  "config_flow": true
+  "config_flow": true,
+  "zeroconf": ["_nut._tcp.local."]
 }

homeassistant/components/nut/translations/en.json

@@ -1,5 +1,6 @@
 {
     "config": {
+        "flow_title": "NUT server ({host}:{port})",
         "abort": {
             "already_configured": "Device is already configured"
         },

homeassistant/generated/zeroconf.py

@@ -40,6 +40,9 @@ ZEROCONF = {
     "_miio._udp.local.": [
         "xiaomi_miio"
     ],
+    "_nut._tcp.local.": [
+        "nut"
+    ],
     "_printer._tcp.local.": [
         "brother"
     ],

tests/components/nut/test_config_flow.py

@@ -1,4 +1,5 @@
 """Test the Network UPS Tools (NUT) config flow."""
+
 from homeassistant import config_entries, data_entry_flow, setup
 from homeassistant.components.nut.const import DOMAIN
 from homeassistant.const import CONF_RESOURCES, CONF_SCAN_INTERVAL
@@ -16,6 +17,59 @@ VALID_CONFIG = {
 }
 
 
+async def test_form_zeroconf(hass):
+    """Test we can setup from zeroconf."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={"source": config_entries.SOURCE_ZEROCONF},
+        data={"host": "192.168.1.5", "port": 1234},
+    )
+    assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
+    assert result["step_id"] == "user"
+    assert result["errors"] == {}
+
+    mock_pynut = _get_mock_pynutclient(
+        list_vars={"battery.voltage": "voltage", "ups.status": "OL"}, list_ups=["ups1"]
+    )
+
+    with patch(
+        "homeassistant.components.nut.PyNUTClient", return_value=mock_pynut,
+    ):
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {"username": "test-username", "password": "test-password"},
+        )
+
+    assert result2["step_id"] == "resources"
+    assert result2["type"] == "form"
+
+    with patch(
+        "homeassistant.components.nut.PyNUTClient", return_value=mock_pynut,
+    ), patch(
+        "homeassistant.components.nut.async_setup", return_value=True
+    ) as mock_setup, patch(
+        "homeassistant.components.nut.async_setup_entry", return_value=True,
+    ) as mock_setup_entry:
+        result3 = await hass.config_entries.flow.async_configure(
+            result2["flow_id"],
+            {"resources": ["battery.voltage", "ups.status", "ups.status.display"]},
+        )
+
+    assert result3["type"] == "create_entry"
+    assert result3["title"] == "192.168.1.5:1234"
+    assert result3["data"] == {
+        "host": "192.168.1.5",
+        "password": "test-password",
+        "port": 1234,
+        "resources": ["battery.voltage", "ups.status", "ups.status.display"],
+        "username": "test-username",
+    }
+    assert result3["result"].unique_id is None
+    await hass.async_block_till_done()
+    assert len(mock_setup.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
 async def test_form_user_one_ups(hass):
     """Test we get the form."""
     await setup.async_setup_component(hass, "persistent_notification", {})