Use eventing for some of the upnp sensors, instead of polling (#120262)

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(

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)

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",

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:

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)

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
 

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: