HEOS confirm discovered devices before adding

homeassistant/components/heos/config_flow.py

@@ -4,9 +4,9 @@
 import voluptuous as vol
 
 from homeassistant import config_entries
-from homeassistant.const import CONF_HOST
+from homeassistant.const import CONF_HOST, CONF_NAME
 
-from .const import DOMAIN
+from .const import DATA_DISCOVERED_HOSTS, DOMAIN
 
 
 def format_title(host: str) -> str:
@@ -23,13 +23,17 @@ class HeosFlowHandler(config_entries.ConfigFlow):
 
     async def async_step_discovery(self, discovery_info):
         """Handle a discovered Heos device."""
-        # Only continue if this is the only active flow
-        flows = self.hass.config_entries.flow.async_progress()
-        heos_flows = [flow for flow in flows if flow['handler'] == DOMAIN]
-        if len(heos_flows) == 1:
-            return await self.async_step_user(
-                {CONF_HOST: discovery_info[CONF_HOST]})
-        return self.async_abort(reason='already_setup')
+        # Store discovered host
+        friendly_name = "{} ({})".format(
+            discovery_info[CONF_NAME], discovery_info[CONF_HOST])
+        self.hass.data.setdefault(DATA_DISCOVERED_HOSTS, {})
+        self.hass.data[DATA_DISCOVERED_HOSTS][friendly_name] \
+            = discovery_info[CONF_HOST]
+        # Abort if other flows in progress or an entry already exists
+        if self._async_in_progress() or self._async_current_entries():
+            return self.async_abort(reason='already_setup')
+        # Show selection form
+        return self.async_show_form(step_id='user')
 
     async def async_step_import(self, user_input=None):
         """Occurs when an entry is setup through config."""
@@ -41,30 +45,33 @@ async def async_step_import(self, user_input=None):
     async def async_step_user(self, user_input=None):
         """Obtain host and validate connection."""
         from pyheos import Heos
-
+        self.hass.data.setdefault(DATA_DISCOVERED_HOSTS, {})
         # Only a single entry is needed for all devices
-        entries = self.hass.config_entries.async_entries(DOMAIN)
-        if entries:
+        if self._async_current_entries():
             return self.async_abort(reason='already_setup')
-
         # Try connecting to host if provided
         errors = {}
         host = None
         if user_input is not None:
             host = user_input[CONF_HOST]
+            # Map host from friendly name if in discovered hosts
+            host = self.hass.data[DATA_DISCOVERED_HOSTS].get(host, host)
             heos = Heos(host)
             try:
                 await heos.connect()
-                return await self.async_step_import(user_input)
+                self.hass.data.pop(DATA_DISCOVERED_HOSTS)
+                return await self.async_step_import({CONF_HOST: host})
             except (asyncio.TimeoutError, ConnectionError):
                 errors[CONF_HOST] = 'connection_failure'
             finally:
                 await heos.disconnect()
 
         # Return form
+        host_type = str if not self.hass.data[DATA_DISCOVERED_HOSTS] \
+            else vol.In(list(self.hass.data[DATA_DISCOVERED_HOSTS]))
         return self.async_show_form(
             step_id='user',
             data_schema=vol.Schema({
-                vol.Required(CONF_HOST, default=host): str
+                vol.Required(CONF_HOST, default=host): host_type
             }),
             errors=errors)

homeassistant/components/heos/const.py

@@ -4,5 +4,6 @@
 COMMAND_RETRY_DELAY = 1
 DATA_CONTROLLER = "controller"
 DATA_SOURCE_MANAGER = "source_manager"
+DATA_DISCOVERED_HOSTS = "heos_discovered_hosts"
 DOMAIN = 'heos'
 SIGNAL_HEOS_SOURCES_UPDATED = "heos_sources_updated"

homeassistant/components/heos/strings.json

@@ -1,12 +1,12 @@
 {
   "config": {
-    "title": "Heos",
+    "title": "HEOS",
     "step": {
       "user": {
         "title": "Connect to Heos",
         "description": "Please enter the host name or IP address of a Heos device (preferably one connected via wire to the network).",
         "data": {
-          "access_token": "Host"
+          "host": "Host"
         }
       }
     },

tests/components/heos/test_config_flow.py

@@ -3,8 +3,8 @@
 
 from homeassistant import data_entry_flow
 from homeassistant.components.heos.config_flow import HeosFlowHandler
-from homeassistant.components.heos.const import DOMAIN
-from homeassistant.const import CONF_HOST
+from homeassistant.components.heos.const import DATA_DISCOVERED_HOSTS, DOMAIN
+from homeassistant.const import CONF_HOST, CONF_NAME
 
 
 async def test_flow_aborts_already_setup(hass, config_entry):
@@ -58,41 +58,51 @@ async def test_create_entry_when_host_valid(hass, controller):
     assert controller.disconnect.call_count == 1
 
 
-async def test_create_entry_with_discovery(hass, controller, discovery_data):
-    """Test discovery creates entry."""
+async def test_create_entry_when_friendly_name_valid(hass, controller):
+    """Test result type is create entry when friendly name is valid."""
+    hass.data[DATA_DISCOVERED_HOSTS] = {"Office (127.0.0.1)": "127.0.0.1"}
+    flow = HeosFlowHandler()
+    flow.hass = hass
+    data = {CONF_HOST: "Office (127.0.0.1)"}
+    result = await flow.async_step_user(data)
+    assert result['type'] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
+    assert result['title'] == 'Controller (127.0.0.1)'
+    assert result['data'] == {CONF_HOST: "127.0.0.1"}
+    assert controller.connect.call_count == 1
+    assert controller.disconnect.call_count == 1
+    assert not hass.data[DATA_DISCOVERED_HOSTS]
+
+
+async def test_discovery_shows_create_form(hass, controller, discovery_data):
+    """Test discovery shows form to confirm setup and subsequent abort."""
     await hass.config_entries.flow.async_init(
                 DOMAIN, context={'source': 'discovery'},
                 data=discovery_data)
     await hass.async_block_till_done()
-    entries = hass.config_entries.async_entries(DOMAIN)
-    assert len(entries) == 1
-    assert entries[0].data == {CONF_HOST: discovery_data[CONF_HOST]}
-    assert entries[0].title == 'Controller (127.0.0.1)'
+    assert len(hass.config_entries.flow.async_progress()) == 1
+    assert hass.data[DATA_DISCOVERED_HOSTS] == {
+        "Office (127.0.0.1)": "127.0.0.1"
+    }
 
-
-async def test_entry_already_exists_discovery(
-        hass, controller, discovery_data, config_entry):
-    """Test discovery does not create multiple entries when already setup."""
-    config_entry.add_to_hass(hass)
+    discovery_data[CONF_HOST] = "127.0.0.2"
+    discovery_data[CONF_NAME] = "Bedroom"
     await hass.config_entries.flow.async_init(
                 DOMAIN, context={'source': 'discovery'},
                 data=discovery_data)
     await hass.async_block_till_done()
-    entries = hass.config_entries.async_entries(DOMAIN)
-    assert len(entries) == 1
-
-
-async def test_multiple_discovery_creates_single_entry(
-        hass, controller, discovery_data):
-    """Test discovery of multiple devices creates a single entry."""
-    hass.async_create_task(
-        hass.config_entries.flow.async_init(
-            DOMAIN, context={'source': 'discovery'},
-            data={CONF_HOST: discovery_data}))
-    hass.async_create_task(
-        hass.config_entries.flow.async_init(
-            DOMAIN, context={'source': 'discovery'},
-            data={CONF_HOST: discovery_data}))
-    await hass.async_block_till_done()
-    entries = hass.config_entries.async_entries(DOMAIN)
-    assert len(entries) == 1
+    assert len(hass.config_entries.flow.async_progress()) == 1
+    assert hass.data[DATA_DISCOVERED_HOSTS] == {
+        "Office (127.0.0.1)": "127.0.0.1",
+        "Bedroom (127.0.0.2)": "127.0.0.2"
+    }
+
+
+async def test_disovery_flow_aborts_already_setup(
+        hass, controller, discovery_data, config_entry):
+    """Test discovery flow aborts when entry already setup."""
+    config_entry.add_to_hass(hass)
+    flow = HeosFlowHandler()
+    flow.hass = hass
+    result = await flow.async_step_discovery(discovery_data)
+    assert result['type'] == data_entry_flow.RESULT_TYPE_ABORT
+    assert result['reason'] == 'already_setup'