From f8b192bd944b442d9af3118027cf6748e0677b79 Mon Sep 17 00:00:00 2001
From: Steven Looman <steven.looman@gmail.com>
Date: Wed, 2 Oct 2024 20:55:25 +0200
Subject: [PATCH] Handle the correct exception type when subscribing to the
 router service returns an error in the upnp component (#127006)

* Catch the right exception when handling subscription errors

* Assert device is forced to poll
---
 homeassistant/components/upnp/device.py |  9 +++--
 tests/components/upnp/test_init.py      | 47 ++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/homeassistant/components/upnp/device.py b/homeassistant/components/upnp/device.py
index 923d4828879..7067d1d2e1a 100644
--- a/homeassistant/components/upnp/device.py
+++ b/homeassistant/components/upnp/device.py
@@ -11,7 +11,7 @@ from urllib.parse import urlparse
 from async_upnp_client.aiohttp import AiohttpNotifyServer, AiohttpSessionRequester
 from async_upnp_client.client_factory import UpnpFactory
 from async_upnp_client.const import AddressTupleVXType
-from async_upnp_client.exceptions import UpnpConnectionError
+from async_upnp_client.exceptions import UpnpCommunicationError
 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
@@ -206,7 +206,7 @@ class Device:
         """Subscribe to services."""
         try:
             await self._igd_device.async_subscribe_services(auto_resubscribe=True)
-        except UpnpConnectionError as ex:
+        except UpnpCommunicationError as ex:
             _LOGGER.debug(
                 "Error subscribing to services, falling back to forced polling: %s", ex
             )
@@ -214,7 +214,10 @@ class Device:
 
     async def async_unsubscribe_services(self) -> None:
         """Unsubscribe from services."""
-        await self._igd_device.async_unsubscribe_services()
+        try:
+            await self._igd_device.async_unsubscribe_services()
+        except UpnpCommunicationError as ex:
+            _LOGGER.debug("Error unsubscribing to services: %s", ex)
 
     async def async_get_data(
         self, entity_description_keys: list[str] | None
diff --git a/tests/components/upnp/test_init.py b/tests/components/upnp/test_init.py
index 0e8551dd8a1..ff74ca87b12 100644
--- a/tests/components/upnp/test_init.py
+++ b/tests/components/upnp/test_init.py
@@ -7,6 +7,7 @@ import copy
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
+from async_upnp_client.exceptions import UpnpCommunicationError
 from async_upnp_client.profiles.igd import IgdDevice
 import pytest
 
@@ -179,7 +180,7 @@ async def test_async_setup_udn_mismatch(
 async def test_async_setup_entry_force_poll(
     hass: HomeAssistant, mock_igd_device: IgdDevice
 ) -> None:
-    """Test async_setup_entry."""
+    """Test async_setup_entry with forced polling."""
     entry = MockConfigEntry(
         domain=DOMAIN,
         unique_id=TEST_USN,
@@ -200,3 +201,47 @@ async def test_async_setup_entry_force_poll(
     assert await hass.config_entries.async_setup(entry.entry_id) is True
 
     mock_igd_device.async_subscribe_services.assert_not_called()
+
+    # Ensure that the device is forced to poll.
+    mock_igd_device.async_get_traffic_and_status_data.assert_called_with(
+        None, force_poll=True
+    )
+
+
+@pytest.mark.usefixtures(
+    "ssdp_instant_discovery",
+    "mock_get_source_ip",
+    "mock_mac_address_from_host",
+)
+async def test_async_setup_entry_force_poll_subscribe_error(
+    hass: HomeAssistant, mock_igd_device: IgdDevice
+) -> None:
+    """Test async_setup_entry where subscribing fails."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        unique_id=TEST_USN,
+        data={
+            CONFIG_ENTRY_ST: TEST_ST,
+            CONFIG_ENTRY_UDN: TEST_UDN,
+            CONFIG_ENTRY_ORIGINAL_UDN: TEST_UDN,
+            CONFIG_ENTRY_LOCATION: TEST_LOCATION,
+            CONFIG_ENTRY_MAC_ADDRESS: TEST_MAC_ADDRESS,
+        },
+        options={
+            CONFIG_ENTRY_FORCE_POLL: False,
+        },
+    )
+
+    # Subscribing partially succeeds, but not completely.
+    # Unsubscribing will fail for the subscribed services afterwards.
+    mock_igd_device.async_subscribe_services.side_effect = UpnpCommunicationError
+    mock_igd_device.async_unsubscribe_services.side_effect = UpnpCommunicationError
+
+    # Load config_entry, should still be able to load, falling back to polling/the old functionality.
+    entry.add_to_hass(hass)
+    assert await hass.config_entries.async_setup(entry.entry_id) is True
+
+    # Ensure that the device is forced to poll.
+    mock_igd_device.async_get_traffic_and_status_data.assert_called_with(
+        None, force_poll=True
+    )