From 1d3fcc67b8fcb60de212e0f33d5e46473766a57a Mon Sep 17 00:00:00 2001 From: Andrew Sayre <6730289+andrewsayre@users.noreply.github.com> Date: Wed, 19 Feb 2025 11:51:47 -0600 Subject: [PATCH] Select preferred discovered HEOS host (#138779) * Select preffered host from discovery * Remove invalid test comment --- homeassistant/components/heos/config_flow.py | 69 ++++++++++------ homeassistant/components/heos/const.py | 1 + homeassistant/components/heos/strings.json | 5 ++ tests/components/heos/__init__.py | 1 + tests/components/heos/test_config_flow.py | 83 ++++++++++---------- tests/components/heos/test_diagnostics.py | 17 ++-- 6 files changed, 101 insertions(+), 75 deletions(-) diff --git a/homeassistant/components/heos/config_flow.py b/homeassistant/components/heos/config_flow.py index db2abee559c..ac09b7ca6bc 100644 --- a/homeassistant/components/heos/config_flow.py +++ b/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 ( - ATTR_UPNP_FRIENDLY_NAME, - SsdpServiceInfo, -) +from homeassistant.helpers.service_info.ssdp import 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 - return self.async_show_form(step_id="user") + # Connect to discovered host and get system information + 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, ) diff --git a/homeassistant/components/heos/const.py b/homeassistant/components/heos/const.py index e9ab51bf16e..6d603f7ad30 100644 --- a/homeassistant/components/heos/const.py +++ b/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" diff --git a/homeassistant/components/heos/strings.json b/homeassistant/components/heos/strings.json index cd3f0b998a1..340eecb9f8b 100644 --- a/homeassistant/components/heos/strings.json +++ b/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%]" diff --git a/tests/components/heos/__init__.py b/tests/components/heos/__init__.py index 0b8aed91edf..5b112f2b986 100644 --- a/tests/components/heos/__init__.py +++ b/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() diff --git a/tests/components/heos/test_config_flow.py b/tests/components/heos/test_config_flow.py index cbc32526958..552b667b6c8 100644 --- a/tests/components/heos/test_config_flow.py +++ b/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( - 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( +async def test_discovery( hass: HomeAssistant, discovery_data: SsdpServiceInfo, discovery_data_bedroom: SsdpServiceInfo, + controller: MockHeos, + system: HeosSystem, ) -> None: - """Test discovery shows form to confirm setup.""" - - # Single discovered host shows form for user to finish setup. - 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. + """Test discovery shows form to confirm, then creates entry.""" + # Single discovered, selects preferred host, shows confirm + controller.get_system_info.return_value = system result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": SOURCE_SSDP}, data=discovery_data_bedroom ) - assert hass.data[DOMAIN] == { - "Office (127.0.0.1)": "127.0.0.1", - "Bedroom (127.0.0.2)": "127.0.0.2", - } - assert result["type"] is FlowResultType.ABORT - assert result["reason"] == "already_in_progress" + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "confirm_discovery" + assert controller.connect.call_count == 1 + assert controller.get_system_info.call_count == 1 + assert controller.disconnect.call_count == 1 + + # 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: diff --git a/tests/components/heos/test_diagnostics.py b/tests/components/heos/test_diagnostics.py index fb71682fb48..a5341ef8d83 100644 --- a/tests/components/heos/test_diagnostics.py +++ b/tests/components/heos/test_diagnostics.py @@ -1,8 +1,6 @@ """Tests for the HEOS diagnostics module.""" -from unittest import mock - -from pyheos import HeosSystem +from pyheos import HeosError, 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, controller.get_system_info.__name__, return_value=system - ): - diagnostics = await get_diagnostics_for_config_entry( - hass, hass_client, config_entry - ) + controller.get_system_info.return_value = system + 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