Select preferred discovered HEOS host (#138779)

* Select preffered host from discovery

* Remove invalid test comment

homeassistant/components/heos/config_flow.py

@@ -17,12 +17,9 @@ from homeassistant.config_entries import (
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import callback
 from homeassistant.helpers import selector
-from homeassistant.helpers.service_info.ssdp import (
+from homeassistant.helpers.service_info.ssdp import SsdpServiceInfo
-    ATTR_UPNP_FRIENDLY_NAME,
-    SsdpServiceInfo,
-)
 
-from .const import DOMAIN
+from .const import DOMAIN, ENTRY_TITLE
 from .coordinator import HeosConfigEntry
 
 _LOGGER = logging.getLogger(__name__)
@@ -37,11 +34,6 @@ AUTH_SCHEMA = vol.Schema(
 )
 
 
-def format_title(host: str) -> str:
-    """Format the title for config entries."""
-    return f"HEOS System (via {host})"
-
-
 async def _validate_host(host: str, errors: dict[str, str]) -> bool:
     """Validate host is reachable, return True, otherwise populate errors and return False."""
     heos = Heos(HeosOptions(host, events=False, heart_beat=False))
@@ -104,6 +96,10 @@ class HeosFlowHandler(ConfigFlow, domain=DOMAIN):
 
     VERSION = 1
 
+    def __init__(self) -> None:
+        """Initialize the HEOS flow."""
+        self._discovered_host: str | None = None
+
     @staticmethod
     @callback
     def async_get_options_flow(config_entry: HeosConfigEntry) -> OptionsFlow:
@@ -117,40 +113,63 @@ class HeosFlowHandler(ConfigFlow, domain=DOMAIN):
         # Store discovered host
         if TYPE_CHECKING:
             assert discovery_info.ssdp_location
-        hostname = urlparse(discovery_info.ssdp_location).hostname
+
-        friendly_name = f"{discovery_info.upnp[ATTR_UPNP_FRIENDLY_NAME]} ({hostname})"
-        self.hass.data.setdefault(DOMAIN, {})
-        self.hass.data[DOMAIN][friendly_name] = hostname
         await self.async_set_unique_id(DOMAIN)
-        # Show selection form
+        # Connect to discovered host and get system information
-        return self.async_show_form(step_id="user")
+        hostname = urlparse(discovery_info.ssdp_location).hostname
+        assert hostname is not None
+        heos = Heos(HeosOptions(hostname, events=False, heart_beat=False))
+        try:
+            await heos.connect()
+            system_info = await heos.get_system_info()
+        except HeosError as error:
+            _LOGGER.debug(
+                "Failed to retrieve system information from discovered HEOS device %s",
+                hostname,
+                exc_info=error,
+            )
+            return self.async_abort(reason="cannot_connect")
+        finally:
+            await heos.disconnect()
+
+        # Select the preferred host, if available
+        if system_info.preferred_hosts:
+            hostname = system_info.preferred_hosts[0].ip_address
+        self._discovered_host = hostname
+        return await self.async_step_confirm_discovery()
+
+    async def async_step_confirm_discovery(
+        self, user_input: dict[str, Any] | None = None
+    ) -> ConfigFlowResult:
+        """Confirm discovered HEOS system."""
+        if user_input is not None:
+            assert self._discovered_host is not None
+            return self.async_create_entry(
+                title=ENTRY_TITLE, data={CONF_HOST: self._discovered_host}
+            )
+
+        self._set_confirm_only()
+        return self.async_show_form(step_id="confirm_discovery")
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Obtain host and validate connection."""
-        self.hass.data.setdefault(DOMAIN, {})
         await self.async_set_unique_id(DOMAIN)
         # Try connecting to host if provided
         errors: dict[str, str] = {}
         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[DOMAIN].get(host, host)
             if await _validate_host(host, errors):
-                self.hass.data.pop(DOMAIN)  # Remove discovery data
                 return self.async_create_entry(
-                    title=format_title(host), data={CONF_HOST: host}
+                    title=ENTRY_TITLE, data={CONF_HOST: host}
                 )
 
         # Return form
-        host_type = (
-            str if not self.hass.data[DOMAIN] else vol.In(list(self.hass.data[DOMAIN]))
-        )
         return self.async_show_form(
             step_id="user",
-            data_schema=vol.Schema({vol.Required(CONF_HOST, default=host): host_type}),
+            data_schema=vol.Schema({vol.Required(CONF_HOST, default=host): str}),
             errors=errors,
         )
 

homeassistant/components/heos/const.py

@@ -3,6 +3,7 @@
 ATTR_PASSWORD = "password"
 ATTR_USERNAME = "username"
 DOMAIN = "heos"
+ENTRY_TITLE = "HEOS System"
 SERVICE_GROUP_VOLUME_SET = "group_volume_set"
 SERVICE_GROUP_VOLUME_DOWN = "group_volume_down"
 SERVICE_GROUP_VOLUME_UP = "group_volume_up"

homeassistant/components/heos/strings.json

@@ -11,6 +11,10 @@
           "host": "Host name or IP address of a HEOS-capable product (preferably one connected via wire to the network)."
         }
       },
+      "confirm_discovery": {
+        "title": "Discovered HEOS System",
+        "description": "Do you want to add your HEOS devices to Home Assistant?"
+      },
       "reconfigure": {
         "title": "Reconfigure HEOS",
         "description": "Change the host name or IP address of the HEOS-capable product used to access your HEOS System.",
@@ -43,6 +47,7 @@
     },
     "abort": {
       "already_in_progress": "[%key:common::config_flow::abort::already_in_progress%]",
+      "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
       "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]",
       "reconfigure_successful": "[%key:common::config_flow::abort::reconfigure_successful%]",
       "single_instance_allowed": "[%key:common::config_flow::abort::single_instance_allowed%]"

tests/components/heos/__init__.py

@@ -22,6 +22,7 @@ class MockHeos(Heos):
         self.get_players: AsyncMock = AsyncMock()
         self.group_volume_down: AsyncMock = AsyncMock()
         self.group_volume_up: AsyncMock = AsyncMock()
+        self.get_system_info: AsyncMock = AsyncMock()
         self.load_players: AsyncMock = AsyncMock()
         self.play_media: AsyncMock = AsyncMock()
         self.play_preset_station: AsyncMock = AsyncMock()

tests/components/heos/test_config_flow.py

@@ -2,7 +2,7 @@
 
 from typing import Any
 
-from pyheos import CommandAuthenticationError, CommandFailedError, HeosError
+from pyheos import CommandAuthenticationError, CommandFailedError, HeosError, HeosSystem
 import pytest
 
 from homeassistant.components.heos.const import DOMAIN
@@ -69,57 +69,46 @@ async def test_create_entry_when_host_valid(
     )
     assert result["type"] is FlowResultType.CREATE_ENTRY
     assert result["result"].unique_id == DOMAIN
-    assert result["title"] == "HEOS System (via 127.0.0.1)"
+    assert result["title"] == "HEOS System"
     assert result["data"] == data
     assert controller.connect.call_count == 2  # Also called in async_setup_entry
     assert controller.disconnect.call_count == 1
 
 
-async def test_create_entry_when_friendly_name_valid(
+async def test_discovery(
-    hass: HomeAssistant, controller: MockHeos
-) -> None:
-    """Test result type is create entry when friendly name is valid."""
-    hass.data[DOMAIN] = {"Office (127.0.0.1)": "127.0.0.1"}
-    data = {CONF_HOST: "Office (127.0.0.1)"}
-
-    result = await hass.config_entries.flow.async_init(
-        DOMAIN, context={"source": SOURCE_USER}, data=data
-    )
-
-    assert result["type"] is FlowResultType.CREATE_ENTRY
-    assert result["result"].unique_id == DOMAIN
-    assert result["title"] == "HEOS System (via 127.0.0.1)"
-    assert result["data"] == {CONF_HOST: "127.0.0.1"}
-    assert controller.connect.call_count == 2  # Also called in async_setup_entry
-    assert controller.disconnect.call_count == 1
-    assert DOMAIN not in hass.data
-
-
-async def test_discovery_shows_create_form(
     hass: HomeAssistant,
     discovery_data: SsdpServiceInfo,
     discovery_data_bedroom: SsdpServiceInfo,
+    controller: MockHeos,
+    system: HeosSystem,
 ) -> None:
-    """Test discovery shows form to confirm setup."""
+    """Test discovery shows form to confirm, then creates entry."""
-
+    # Single discovered, selects preferred host, shows confirm
-    # Single discovered host shows form for user to finish setup.
+    controller.get_system_info.return_value = system
-    result = await hass.config_entries.flow.async_init(
-        DOMAIN, context={"source": SOURCE_SSDP}, data=discovery_data
-    )
-    assert hass.data[DOMAIN] == {"Office (127.0.0.1)": "127.0.0.1"}
-    assert result["type"] is FlowResultType.FORM
-    assert result["step_id"] == "user"
-
-    # Subsequent discovered hosts append to discovered hosts and abort.
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": SOURCE_SSDP}, data=discovery_data_bedroom
     )
-    assert hass.data[DOMAIN] == {
+    assert result["type"] is FlowResultType.FORM
-        "Office (127.0.0.1)": "127.0.0.1",
+    assert result["step_id"] == "confirm_discovery"
-        "Bedroom (127.0.0.2)": "127.0.0.2",
+    assert controller.connect.call_count == 1
-    }
+    assert controller.get_system_info.call_count == 1
-    assert result["type"] is FlowResultType.ABORT
+    assert controller.disconnect.call_count == 1
-    assert result["reason"] == "already_in_progress"
+
+    # Subsequent discovered hosts abort.
+    subsequent_result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": SOURCE_SSDP}, data=discovery_data
+    )
+    assert subsequent_result["type"] is FlowResultType.ABORT
+    assert subsequent_result["reason"] == "already_in_progress"
+
+    # Confirm set up
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], user_input={}
+    )
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert result["result"].unique_id == DOMAIN
+    assert result["title"] == "HEOS System"
+    assert result["data"] == {CONF_HOST: "127.0.0.1"}
 
 
 async def test_discovery_flow_aborts_already_setup(
@@ -136,6 +125,20 @@ async def test_discovery_flow_aborts_already_setup(
     assert result["reason"] == "single_instance_allowed"
 
 
+async def test_discovery_fails_to_connect_aborts(
+    hass: HomeAssistant, discovery_data: SsdpServiceInfo, controller: MockHeos
+) -> None:
+    """Test discovery aborts when trying to connect to host."""
+    controller.connect.side_effect = HeosError()
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": SOURCE_SSDP}, data=discovery_data
+    )
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "cannot_connect"
+    assert controller.connect.call_count == 1
+    assert controller.disconnect.call_count == 1
+
+
 async def test_reconfigure_validates_and_updates_config(
     hass: HomeAssistant, config_entry: MockConfigEntry, controller: MockHeos
 ) -> None:

tests/components/heos/test_diagnostics.py

@@ -1,8 +1,6 @@
 """Tests for the HEOS diagnostics module."""
 
-from unittest import mock
+from pyheos import HeosError, HeosSystem
-
-from pyheos import HeosSystem
 import pytest
 from syrupy.assertion import SnapshotAssertion
 from syrupy.filters import props
@@ -33,12 +31,10 @@ async def test_config_entry_diagnostics(
     config_entry.add_to_hass(hass)
     assert await hass.config_entries.async_setup(config_entry.entry_id)
 
-    with mock.patch.object(
+    controller.get_system_info.return_value = system
-        controller, controller.get_system_info.__name__, return_value=system
+    diagnostics = await get_diagnostics_for_config_entry(
-    ):
+        hass, hass_client, config_entry
-        diagnostics = await get_diagnostics_for_config_entry(
+    )
-            hass, hass_client, config_entry
-        )
 
     assert diagnostics == snapshot(
         exclude=props("created_at", "modified_at", "entry_id")
@@ -50,13 +46,14 @@ async def test_config_entry_diagnostics_error_getting_system(
     hass: HomeAssistant,
     hass_client: ClientSessionGenerator,
     config_entry: MockConfigEntry,
+    controller: MockHeos,
     snapshot: SnapshotAssertion,
 ) -> None:
     """Test generating diagnostics with error during getting system info."""
     config_entry.add_to_hass(hass)
     assert await hass.config_entries.async_setup(config_entry.entry_id)
 
-    # Not patching get_system_info to raise error 'Not connected to device'
+    controller.get_system_info.side_effect = HeosError("Not connected to device")
 
     diagnostics = await get_diagnostics_for_config_entry(
         hass, hass_client, config_entry