From 7a098cff1ca64c41a79fabce699ac147982ffb06 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Wed, 1 Dec 2021 20:23:38 +0100 Subject: [PATCH] Use dataclass properties in upnp discovery (#60744) Co-authored-by: epenet --- homeassistant/components/upnp/config_flow.py | 29 +++++++++++--------- tests/components/upnp/conftest.py | 4 ++- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/homeassistant/components/upnp/config_flow.py b/homeassistant/components/upnp/config_flow.py index 58dee0c7021..fd4ed9d4051 100644 --- a/homeassistant/components/upnp/config_flow.py +++ b/homeassistant/components/upnp/config_flow.py @@ -13,6 +13,7 @@ from homeassistant.components import ssdp from homeassistant.components.ssdp import SsdpChange from homeassistant.const import CONF_SCAN_INTERVAL from homeassistant.core import HomeAssistant, callback +from homeassistant.data_entry_flow import FlowResult from .const import ( CONFIG_ENTRY_HOSTNAME, @@ -28,22 +29,22 @@ from .const import ( ) -def _friendly_name_from_discovery(discovery_info: Mapping[str, Any]) -> str: +def _friendly_name_from_discovery(discovery_info: ssdp.SsdpServiceInfo) -> str: """Extract user-friendly name from discovery.""" return ( - discovery_info.get("friendlyName") - or discovery_info.get("modeName") - or discovery_info.get("_host", "") + discovery_info.upnp.get(ssdp.ATTR_UPNP_FRIENDLY_NAME) + or discovery_info.upnp.get(ssdp.ATTR_UPNP_MODEL_NAME) + or discovery_info.ssdp_headers.get("_host", "") ) -def _is_complete_discovery(discovery_info: Mapping[str, Any]) -> bool: +def _is_complete_discovery(discovery_info: ssdp.SsdpServiceInfo) -> bool: """Test if discovery is complete and usable.""" return ( - ssdp.ATTR_UPNP_UDN in discovery_info - and discovery_info.get(ssdp.ATTR_SSDP_ST) - and discovery_info.get(ssdp.ATTR_SSDP_LOCATION) - and discovery_info.get(ssdp.ATTR_SSDP_USN) + ssdp.ATTR_UPNP_UDN in discovery_info.upnp + and discovery_info.ssdp_st + and discovery_info.ssdp_location + and discovery_info.ssdp_usn ) @@ -90,7 +91,9 @@ async def _async_wait_for_discoveries(hass: HomeAssistant) -> bool: return True -async def _async_discover_igd_devices(hass: HomeAssistant) -> list[Mapping[str, Any]]: +async def _async_discover_igd_devices( + hass: HomeAssistant, +) -> list[ssdp.SsdpServiceInfo]: """Discovery IGD devices.""" return await ssdp.async_get_discovery_info_by_st( hass, ST_IGD_V1 @@ -206,7 +209,7 @@ class UpnpFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): return await self._async_create_entry_from_discovery(discovery) - async def async_step_ssdp(self, discovery_info: Mapping) -> Mapping[str, Any]: + async def async_step_ssdp(self, discovery_info: ssdp.SsdpServiceInfo) -> FlowResult: """Handle a discovered UPnP/IGD device. This flow is triggered by the SSDP component. It will check if the @@ -220,9 +223,9 @@ class UpnpFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): return self.async_abort(reason="incomplete_discovery") # Ensure not already configuring/configured. - unique_id = discovery_info[ssdp.ATTR_SSDP_USN] + unique_id = discovery_info.ssdp_usn await self.async_set_unique_id(unique_id) - hostname = discovery_info["_host"] + hostname = discovery_info.ssdp_headers["_host"] self._abort_if_unique_id_configured(updates={CONFIG_ENTRY_HOSTNAME: hostname}) # Handle devices changing their UDN, only allow a single host. diff --git a/tests/components/upnp/conftest.py b/tests/components/upnp/conftest.py index 6f6f2905869..d2b1a849ba5 100644 --- a/tests/components/upnp/conftest.py +++ b/tests/components/upnp/conftest.py @@ -38,10 +38,12 @@ TEST_DISCOVERY = ssdp.SsdpServiceInfo( ssdp.ATTR_UPNP_UDN: TEST_UDN, "usn": TEST_USN, "location": TEST_LOCATION, - "_host": TEST_HOSTNAME, "_udn": TEST_UDN, "friendlyName": TEST_FRIENDLY_NAME, }, + ssdp_headers={ + "_host": TEST_HOSTNAME, + }, )