diff --git a/homeassistant/components/upnp/__init__.py b/homeassistant/components/upnp/__init__.py index ea9930f047f..df8c6326e10 100644 --- a/homeassistant/components/upnp/__init__.py +++ b/homeassistant/components/upnp/__init__.py @@ -82,14 +82,22 @@ async def async_setup_entry(hass: HomeAssistant, entry: UpnpConfigEntry) -> bool assert discovery_info is not None assert discovery_info.ssdp_udn assert discovery_info.ssdp_all_locations + force_poll = False location = get_preferred_location(discovery_info.ssdp_all_locations) try: - device = await async_create_device(hass, location) + device = await async_create_device(hass, location, force_poll) except UpnpConnectionError as err: raise ConfigEntryNotReady( f"Error connecting to device at location: {location}, err: {err}" ) from err + # Try to subscribe, if configured. + if not force_poll: + await device.async_subscribe_services() + + # Unsubscribe services on unload. + entry.async_on_unload(device.async_unsubscribe_services) + # Track the original UDN such that existing sensors do not change their unique_id. if CONFIG_ENTRY_ORIGINAL_UDN not in entry.data: hass.config_entries.async_update_entry( diff --git a/homeassistant/components/upnp/binary_sensor.py b/homeassistant/components/upnp/binary_sensor.py index 9784f9c6e0b..fb32946bf7d 100644 --- a/homeassistant/components/upnp/binary_sensor.py +++ b/homeassistant/components/upnp/binary_sensor.py @@ -51,8 +51,8 @@ async def async_setup_entry( for entity_description in SENSOR_DESCRIPTIONS if coordinator.data.get(entity_description.key) is not None ] - LOGGER.debug("Adding binary_sensor entities: %s", entities) async_add_entities(entities) + LOGGER.debug("Added binary_sensor entities: %s", entities) class UpnpStatusBinarySensor(UpnpEntity, BinarySensorEntity): @@ -72,3 +72,13 @@ class UpnpStatusBinarySensor(UpnpEntity, BinarySensorEntity): def is_on(self) -> bool: """Return true if the binary sensor is on.""" return self.coordinator.data[self.entity_description.key] == "Connected" + + async def async_added_to_hass(self) -> None: + """Subscribe to updates.""" + await super().async_added_to_hass() + + # Register self at coordinator. + key = self.entity_description.key + entity_id = self.entity_id + unregister = self.coordinator.register_entity(key, entity_id) + self.async_on_remove(unregister) diff --git a/homeassistant/components/upnp/coordinator.py b/homeassistant/components/upnp/coordinator.py index 72e14ecc4ff..a4cb608615c 100644 --- a/homeassistant/components/upnp/coordinator.py +++ b/homeassistant/components/upnp/coordinator.py @@ -1,5 +1,7 @@ """UPnP/IGD coordinator.""" +from collections import defaultdict +from collections.abc import Callable from datetime import datetime, timedelta from async_upnp_client.exceptions import UpnpCommunicationError @@ -27,6 +29,7 @@ class UpnpDataUpdateCoordinator( """Initialize.""" self.device = device self.device_entry = device_entry + self._features_by_entity_id: defaultdict[str, set[str]] = defaultdict(set) super().__init__( hass, @@ -35,12 +38,35 @@ class UpnpDataUpdateCoordinator( update_interval=update_interval, ) + def register_entity(self, key: str, entity_id: str) -> Callable[[], None]: + """Register an entity.""" + # self._entities.append(entity) + self._features_by_entity_id[key].add(entity_id) + + def unregister_entity() -> None: + """Unregister entity.""" + self._features_by_entity_id[key].remove(entity_id) + + if not self._features_by_entity_id[key]: + del self._features_by_entity_id[key] + + return unregister_entity + + @property + def _entity_description_keys(self) -> list[str] | None: + """Return a list of entity description keys for which data is required.""" + if not self._features_by_entity_id: + # Must be the first update, no entities attached/enabled yet. + return None + + return list(self._features_by_entity_id.keys()) + async def _async_update_data( self, ) -> dict[str, str | datetime | int | float | None]: """Update data.""" try: - return await self.device.async_get_data() + return await self.device.async_get_data(self._entity_description_keys) except UpnpCommunicationError as exception: LOGGER.debug( "Caught exception when updating device: %s, exception: %s", diff --git a/homeassistant/components/upnp/device.py b/homeassistant/components/upnp/device.py index bb0bcfc6a6e..e819a16f2d2 100644 --- a/homeassistant/components/upnp/device.py +++ b/homeassistant/components/upnp/device.py @@ -8,9 +8,12 @@ from ipaddress import ip_address from typing import Any from urllib.parse import urlparse -from async_upnp_client.aiohttp import AiohttpSessionRequester +from async_upnp_client.aiohttp import AiohttpNotifyServer, AiohttpSessionRequester from async_upnp_client.client_factory import UpnpFactory -from async_upnp_client.profiles.igd import IgdDevice +from async_upnp_client.const import AddressTupleVXType +from async_upnp_client.exceptions import UpnpConnectionError +from async_upnp_client.profiles.igd import IgdDevice, IgdStateItem +from async_upnp_client.utils import async_get_local_ip from getmac import get_mac_address from homeassistant.core import HomeAssistant @@ -33,6 +36,20 @@ from .const import ( WAN_STATUS, ) +TYPE_STATE_ITEM_MAPPING = { + BYTES_RECEIVED: IgdStateItem.BYTES_RECEIVED, + BYTES_SENT: IgdStateItem.BYTES_SENT, + KIBIBYTES_PER_SEC_RECEIVED: IgdStateItem.KIBIBYTES_PER_SEC_RECEIVED, + KIBIBYTES_PER_SEC_SENT: IgdStateItem.KIBIBYTES_PER_SEC_SENT, + PACKETS_PER_SEC_RECEIVED: IgdStateItem.PACKETS_PER_SEC_RECEIVED, + PACKETS_PER_SEC_SENT: IgdStateItem.PACKETS_PER_SEC_SENT, + PACKETS_RECEIVED: IgdStateItem.PACKETS_RECEIVED, + PACKETS_SENT: IgdStateItem.PACKETS_SENT, + ROUTER_IP: IgdStateItem.EXTERNAL_IP_ADDRESS, + ROUTER_UPTIME: IgdStateItem.UPTIME, + WAN_STATUS: IgdStateItem.CONNECTION_STATUS, +} + def get_preferred_location(locations: set[str]) -> str: """Get the preferred location (an IPv4 location) from a set of locations.""" @@ -64,26 +81,43 @@ async def async_get_mac_address_from_host(hass: HomeAssistant, host: str) -> str return mac_address -async def async_create_device(hass: HomeAssistant, location: str) -> Device: +async def async_create_device( + hass: HomeAssistant, location: str, force_poll: bool +) -> Device: """Create UPnP/IGD device.""" session = async_get_clientsession(hass, verify_ssl=False) requester = AiohttpSessionRequester(session, with_sleep=True, timeout=20) + # Create UPnP device. factory = UpnpFactory(requester, non_strict=True) upnp_device = await factory.async_create_device(location) + # Create notify server. + _, local_ip = await async_get_local_ip(location) + source: AddressTupleVXType = (local_ip, 0) + notify_server = AiohttpNotifyServer( + requester=requester, + source=source, + ) + await notify_server.async_start_server() + _LOGGER.debug("Started event handler at %s", notify_server.callback_url) + # Create profile wrapper. - igd_device = IgdDevice(upnp_device, None) - return Device(hass, igd_device) + igd_device = IgdDevice(upnp_device, notify_server.event_handler) + return Device(hass, igd_device, force_poll) class Device: """Home Assistant representation of a UPnP/IGD device.""" - def __init__(self, hass: HomeAssistant, igd_device: IgdDevice) -> None: + def __init__( + self, hass: HomeAssistant, igd_device: IgdDevice, force_poll: bool + ) -> None: """Initialize UPnP/IGD device.""" self.hass = hass self._igd_device = igd_device + self._force_poll = force_poll + self.coordinator: ( DataUpdateCoordinator[dict[str, str | datetime | int | float | None]] | None ) = None @@ -151,11 +185,54 @@ class Device: """Get string representation.""" return f"IGD Device: {self.name}/{self.udn}::{self.device_type}" - async def async_get_data(self) -> dict[str, str | datetime | int | float | None]: + @property + def force_poll(self) -> bool: + """Get force_poll.""" + return self._force_poll + + async def async_set_force_poll(self, force_poll: bool) -> None: + """Set force_poll, and (un)subscribe if needed.""" + self._force_poll = force_poll + + if self._force_poll: + # No need for subscriptions, as eventing will never be used. + await self.async_unsubscribe_services() + elif not self._force_poll and not self._igd_device.is_subscribed: + await self.async_subscribe_services() + + async def async_subscribe_services(self) -> None: + """Subscribe to services.""" + try: + await self._igd_device.async_subscribe_services(auto_resubscribe=True) + except UpnpConnectionError as ex: + _LOGGER.debug( + "Error subscribing to services, falling back to forced polling: %s", ex + ) + await self.async_set_force_poll(True) + + async def async_unsubscribe_services(self) -> None: + """Unsubscribe from services.""" + await self._igd_device.async_unsubscribe_services() + + async def async_get_data( + self, entity_description_keys: list[str] | None + ) -> dict[str, str | datetime | int | float | None]: """Get all data from device.""" - _LOGGER.debug("Getting data for device: %s", self) + if not entity_description_keys: + igd_state_items = None + else: + igd_state_items = { + TYPE_STATE_ITEM_MAPPING[key] for key in entity_description_keys + } + + _LOGGER.debug( + "Getting data for device: %s, state_items: %s, force_poll: %s", + self, + igd_state_items, + self._force_poll, + ) igd_state = await self._igd_device.async_get_traffic_and_status_data( - force_poll=True + igd_state_items, force_poll=self._force_poll ) def get_value(value: Any) -> Any: diff --git a/homeassistant/components/upnp/sensor.py b/homeassistant/components/upnp/sensor.py index df7128830b3..d72dce55eab 100644 --- a/homeassistant/components/upnp/sensor.py +++ b/homeassistant/components/upnp/sensor.py @@ -159,8 +159,8 @@ async def async_setup_entry( if coordinator.data.get(entity_description.key) is not None ] - LOGGER.debug("Adding sensor entities: %s", entities) async_add_entities(entities) + LOGGER.debug("Added sensor entities: %s", entities) class UpnpSensor(UpnpEntity, SensorEntity): @@ -174,3 +174,13 @@ class UpnpSensor(UpnpEntity, SensorEntity): if (key := self.entity_description.value_key) is None: return None return self.coordinator.data[key] + + async def async_added_to_hass(self) -> None: + """Subscribe to updates.""" + await super().async_added_to_hass() + + # Register self at coordinator. + key = self.entity_description.key + entity_id = self.entity_id + unregister = self.coordinator.register_entity(key, entity_id) + self.async_on_remove(unregister) diff --git a/tests/components/upnp/conftest.py b/tests/components/upnp/conftest.py index 0bfcd062ac0..8f5de1fa824 100644 --- a/tests/components/upnp/conftest.py +++ b/tests/components/upnp/conftest.py @@ -4,9 +4,11 @@ from __future__ import annotations import copy from datetime import datetime +import socket from unittest.mock import AsyncMock, MagicMock, PropertyMock, create_autospec, patch from urllib.parse import urlparse +from async_upnp_client.aiohttp import AiohttpNotifyServer from async_upnp_client.client import UpnpDevice from async_upnp_client.profiles.igd import IgdDevice, IgdState import pytest @@ -98,9 +100,24 @@ def mock_igd_device(mock_async_create_device) -> IgdDevice: port_mapping_number_of_entries=0, ) - with patch( - "homeassistant.components.upnp.device.IgdDevice.__new__", - return_value=mock_igd_device, + mock_igd_device.async_subscribe_services = AsyncMock() + + mock_notify_server = create_autospec(AiohttpNotifyServer) + mock_notify_server.event_handler = MagicMock() + + with ( + patch( + "homeassistant.components.upnp.device.async_get_local_ip", + return_value=(socket.AF_INET, "127.0.0.1"), + ), + patch( + "homeassistant.components.upnp.device.IgdDevice.__new__", + return_value=mock_igd_device, + ), + patch( + "homeassistant.components.upnp.device.AiohttpNotifyServer.__new__", + return_value=mock_notify_server, + ), ): yield mock_igd_device diff --git a/tests/components/upnp/test_init.py b/tests/components/upnp/test_init.py index 4b5e375f8e0..422d8c9e33a 100644 --- a/tests/components/upnp/test_init.py +++ b/tests/components/upnp/test_init.py @@ -5,6 +5,7 @@ from __future__ import annotations import copy from unittest.mock import AsyncMock, MagicMock, patch +from async_upnp_client.profiles.igd import IgdDevice import pytest from homeassistant.components import ssdp @@ -31,7 +32,9 @@ from tests.common import MockConfigEntry @pytest.mark.usefixtures("ssdp_instant_discovery", "mock_mac_address_from_host") -async def test_async_setup_entry_default(hass: HomeAssistant) -> None: +async def test_async_setup_entry_default( + hass: HomeAssistant, mock_igd_device: IgdDevice +) -> None: """Test async_setup_entry.""" entry = MockConfigEntry( domain=DOMAIN, @@ -49,6 +52,8 @@ async def test_async_setup_entry_default(hass: HomeAssistant) -> None: entry.add_to_hass(hass) assert await hass.config_entries.async_setup(entry.entry_id) is True + mock_igd_device.async_subscribe_services.assert_called() + @pytest.mark.usefixtures("ssdp_instant_discovery", "mock_no_mac_address_from_host") async def test_async_setup_entry_default_no_mac_address(hass: HomeAssistant) -> None: