Merge pull request #982 from plugwise/coord-restore-state

bouwew · web-flow · commit 3817e52a6d9e · 2025-12-05T08:01:20.000+01:00
Coordinator: add _async_setup() and collect set of existing devices
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,7 @@ Versions from 0.40 and up
 
 ## Ongoing
 
+- Improve automatic deletion of device(s) removed from the backend output, via PR [$982](https://github.com/plugwise/plugwise-beta/pull/982)
 - Extended feature: extent/improve Plugwise groups, via PR[#976](https://github.com/plugwise/plugwise-beta/pull/976) and plugwise [v1.11.0](https://github.com/plugwise/python-plugwise/releases/tag/v1.11.0)
 - DeviceInfo: show configuration_url on gateway only, via PR [#975](https://github.com/plugwise/plugwise-beta/pull/975)
 
diff --git a/custom_components/plugwise/coordinator.py b/custom_components/plugwise/coordinator.py
@@ -6,6 +6,7 @@
 from plugwise.exceptions import (
     ConnectionFailedError,
     InvalidAuthentication,
+    InvalidSetupError,
     InvalidXMLError,
     PlugwiseError,
     ResponseError,
@@ -37,6 +38,9 @@ class PlugwiseDataUpdateCoordinator(DataUpdateCoordinator[dict[str, GwEntityData
     """Class to manage fetching Plugwise data from single endpoint."""
 
     _connected: bool = False
+    _current_devices: set[str]
+    _stored_devices: set[str]
+    new_devices: set[str]
 
     config_entry: PlugwiseConfigEntry
 
@@ -73,31 +77,18 @@ def __init__(
             username=self.config_entry.data[CONF_USERNAME],
             websession=async_get_clientsession(hass, verify_ssl=False),
         )
-        self._current_devices: set[str] = set()
-        self.new_devices: set[str] = set()
+        self._current_devices = set()
+        self._stored_devices = set()
+        self.new_devices = set()
         self.update_interval = update_interval
 
     async def _connect(self) -> None:
-        """Connect to the Plugwise Smile."""
-        version = await self.api.connect()
-        self._connected = isinstance(version, Version)
-        if self._connected:
-            self.update_interval = DEFAULT_SCAN_INTERVAL.get(
-                self.api.smile.type, timedelta(seconds=60)
-            )  # pw-beta options scan-interval
-            if (custom_time := self.config_entry.options.get(CONF_SCAN_INTERVAL)) is not None:
-                self.update_interval = timedelta(
-                    seconds=int(custom_time)
-                )  # pragma: no cover  # pw-beta options
-
-        LOGGER.debug("DUC update interval: %s", self.update_interval)  # pw-beta options
+        """Connect to the Plugwise Smile.
 
-    async def _async_update_data(self) -> dict[str, GwEntityData]:
-        """Fetch data from Plugwise."""
+        A Version object is received when the connection succeeds.
+        """
         try:
-            if not self._connected:
-                await self._connect()
-            data = await self.api.async_update()
+            version = await self.api.connect()
         except ConnectionFailedError as err:
             raise UpdateFailed(
                 translation_domain=DOMAIN,
@@ -108,35 +99,73 @@ async def _async_update_data(self) -> dict[str, GwEntityData]:
                 translation_domain=DOMAIN,
                 translation_key="authentication_failed",
             ) from err
-        except (InvalidXMLError, ResponseError) as err:
-            # pwbeta TODO; we had {err} in the text, but not upstream, do we want this?
-            raise UpdateFailed(
+        except InvalidSetupError as err:
+            raise ConfigEntryError(
                 translation_domain=DOMAIN,
-                translation_key="invalid_xml_data",
+                translation_key="invalid_setup",
             ) from err
-        except PlugwiseError as err:
+        except (InvalidXMLError, ResponseError) as err:
             raise UpdateFailed(
                 translation_domain=DOMAIN,
-                translation_key="data_incomplete_or_missing",
+                translation_key="invalid_xml_data",
             ) from err
         except UnsupportedDeviceError as err:
             raise ConfigEntryError(
                 translation_domain=DOMAIN,
                 translation_key="unsupported_firmware",
             ) from err
 
-        LOGGER.debug(f"{self.api.smile.name} data: %s", data)
+        self._connected = isinstance(version, Version)
+        if self._connected:
+            self.update_interval = DEFAULT_SCAN_INTERVAL.get(
+                self.api.smile.type, timedelta(seconds=60)
+            )  # pw-beta options scan-interval
+            if (custom_time := self.config_entry.options.get(CONF_SCAN_INTERVAL)) is not None:
+                self.update_interval = timedelta(
+                    seconds=int(custom_time)
+                )  # pragma: no cover  # pw-beta options
+
+        LOGGER.debug("DUC update interval: %s", self.update_interval)  # pw-beta options
+
+    async def _async_setup(self) -> None:
+        """Initialize the update_data process."""
+        device_reg = dr.async_get(self.hass)
+        device_entries = dr.async_entries_for_config_entry(
+            device_reg, self.config_entry.entry_id
+        )
+        self._stored_devices = {
+            identifier[1]
+            for device_entry in device_entries
+            for identifier in device_entry.identifiers
+            if identifier[0] == DOMAIN
+        }
+
+    async def _async_update_data(self) -> dict[str, GwEntityData]:
+        """Fetch data from Plugwise."""
+        if not self._connected:
+            await self._connect()
+        try:
+            data = await self.api.async_update()
+        except PlugwiseError as err:
+            raise UpdateFailed(
+                translation_domain=DOMAIN,
+                translation_key="data_incomplete_or_missing",
+            ) from err
+
         await self._async_add_remove_devices(data)
+        LOGGER.debug("%s data: %s", self.api.smile.name, data)
         return data
 
     async def _async_add_remove_devices(self, data: dict[str, GwEntityData]) -> None:
         """Add new Plugwise devices, remove non-existing devices."""
-        # Check for new or removed devices
-        self.new_devices = set(data) - self._current_devices
-        removed_devices = self._current_devices - set(data)
-        self._current_devices = set(data)
-
-        if removed_devices:
+        set_of_data = set(data)
+        # Check for new or removed devices,
+        # 'new_devices' contains all devices present in 'data' at init ('self._current_devices' is empty)
+        # this is required for the proper initialization of all the present platform entities.
+        self.new_devices = set_of_data - self._current_devices
+        current_devices = self._stored_devices if not self._current_devices else self._current_devices
+        self._current_devices = set_of_data
+        if (current_devices - set_of_data):  # device(s) to remove
             await self._async_remove_devices(data)
 
     async def _async_remove_devices(self, data: dict[str, GwEntityData]) -> None:
@@ -148,9 +177,10 @@ async def _async_remove_devices(self, data: dict[str, GwEntityData]) -> None:
 
         # First find the Plugwise via_device
         gateway_device = device_reg.async_get_device({(DOMAIN, self.api.gateway_id)})
-        if gateway_device is not None:
-            via_device_id = gateway_device.id
+        if gateway_device is None:
+            return  # pragma: no cover
 
+        via_device_id = gateway_device.id
         # Then remove the connected orphaned device(s)
         for device_entry in device_list:
             for identifier in device_entry.identifiers:
diff --git a/tests/components/plugwise/test_init.py b/tests/components/plugwise/test_init.py
@@ -1,11 +1,11 @@
 """Tests for the Plugwise Climate integration."""
 from datetime import timedelta
-import logging
 from unittest.mock import MagicMock, patch
 
 from plugwise.exceptions import (
     ConnectionFailedError,
     InvalidAuthentication,
+    InvalidSetupError,
     InvalidXMLError,
     PlugwiseError,
     ResponseError,
@@ -15,6 +15,7 @@
 
 from freezegun.api import FrozenDateTimeFactory
 from homeassistant.components.plugwise.const import DOMAIN
+from homeassistant.components.plugwise.coordinator import PlugwiseDataUpdateCoordinator
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import (
     CONF_HOST,
@@ -26,13 +27,13 @@
     Platform,
 )
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryError
 from homeassistant.helpers import device_registry as dr, entity_registry as er
+from homeassistant.helpers.update_coordinator import UpdateFailed
 from homeassistant.setup import async_setup_component
 
 from tests.common import MockConfigEntry, async_fire_time_changed
 
-LOGGER = logging.getLogger(__package__)
-
 HA_PLUGWISE_SMILE_ASYNC_UPDATE = (
     "homeassistant.components.plugwise.coordinator.Smile.async_update"
 )
@@ -96,23 +97,12 @@ async def test_load_unload_config_entry(
 
 @pytest.mark.parametrize("chosen_env", ["anna_heatpump_heating"], indirect=True)
 @pytest.mark.parametrize("cooling_present", [True], indirect=True)
-@pytest.mark.parametrize(
-    ("side_effect", "entry_state"),
-    [
-        (ConnectionFailedError, ConfigEntryState.SETUP_RETRY),
-        (InvalidAuthentication, ConfigEntryState.SETUP_ERROR),
-        (InvalidXMLError, ConfigEntryState.SETUP_RETRY),
-        (ResponseError, ConfigEntryState.SETUP_RETRY),
-        (PlugwiseError, ConfigEntryState.SETUP_RETRY),
-        (UnsupportedDeviceError, ConfigEntryState.SETUP_ERROR),
-    ],
-)
+@pytest.mark.parametrize("side_effect", [PlugwiseError])
 async def test_gateway_config_entry_not_ready(
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,
     mock_smile_anna: MagicMock,
     side_effect: Exception,
-    entry_state: ConfigEntryState,
 ) -> None:
     """Test the Plugwise configuration entry not ready."""
     mock_smile_anna.async_update.side_effect = side_effect
@@ -122,7 +112,38 @@ async def test_gateway_config_entry_not_ready(
     await hass.async_block_till_done()
 
     assert len(mock_smile_anna.connect.mock_calls) == 1
-    assert mock_config_entry.state is entry_state
+    assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+@pytest.mark.parametrize("chosen_env", ["anna_heatpump_heating"], indirect=True)
+@pytest.mark.parametrize("cooling_present", [True], indirect=True)
+@pytest.mark.parametrize(
+    ("side_effect", "expected_raise"),
+    [
+        (ConnectionFailedError, UpdateFailed),
+        (InvalidAuthentication, ConfigEntryError),
+        (InvalidSetupError, ConfigEntryError),
+        (InvalidXMLError, UpdateFailed),
+        (ResponseError, UpdateFailed),
+        (UnsupportedDeviceError, ConfigEntryError),
+    ],
+)
+async def test_coordinator_connect_exceptions(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_smile_anna: MagicMock,
+    side_effect: Exception,
+    expected_raise: Exception,
+) -> None:
+    """Ensure _connect raises translated errors."""
+    mock_smile_anna.connect.side_effect = side_effect
+    coordinator = PlugwiseDataUpdateCoordinator(
+        hass,
+        cooldown=0,
+        config_entry=mock_config_entry,
+    )
+    with pytest.raises(expected_raise):
+        await coordinator._connect()
 
 
 @pytest.mark.parametrize("chosen_env", ["p1v4_442_single"], indirect=True)
@@ -160,7 +181,7 @@ async def check_migration(
     mock_config_entry.add_to_hass(hass)
 
     entity_registry = er.async_get(hass)
-    entity: entity_registry.RegistryEntry = entity_registry.async_get_or_create(
+    entity: er.RegistryEntry = entity_registry.async_get_or_create(
         **entitydata,
         config_entry=mock_config_entry,
     )
@@ -245,38 +266,6 @@ async def test_migrate_unique_id_relay(
         hass, mock_config_entry, entitydata, old_unique_id, new_unique_id
     )
 
-#### pw-beta only ####
-@pytest.mark.parametrize("chosen_env", ["m_anna_heatpump_cooling"], indirect=True)
-@pytest.mark.parametrize("cooling_present", [True], indirect=True)
-async def test_entry_migration(
-    hass: HomeAssistant, mock_smile_anna: MagicMock
-) -> None:
-    """Test config entry 1.2 -> 1.1 migration."""
-    entry = MockConfigEntry(
-        domain=DOMAIN,
-        data={
-            CONF_HOST: "127.0.0.1",
-            CONF_MAC: "AA:BB:CC:DD:EE:FF",
-            CONF_PASSWORD: "test-password",
-            CONF_PORT: 80,
-            CONF_TIMEOUT: 30,
-            CONF_USERNAME: "smile",
-        },
-        minor_version=2,
-        version=1,
-        unique_id="smile98765",
-    )
-
-    entry.runtime_data = MagicMock(api=mock_smile_anna)
-    entry.add_to_hass(hass)
-    await hass.config_entries.async_setup(entry.entry_id)
-    await hass.async_block_till_done()
-
-    assert entry.version == 1
-    assert entry.minor_version == 1
-    assert entry.data.get(CONF_TIMEOUT) is None
-    assert entry.state is ConfigEntryState.LOADED
-
 
 @pytest.mark.parametrize("chosen_env", ["m_adam_heating"], indirect=True)
 @pytest.mark.parametrize("cooling_present", [False], indirect=True)
@@ -358,3 +347,67 @@ async def test_update_device(
         for device_entry in device_registry.devices.values():
             item_list.extend(x[1] for x in device_entry.identifiers)
         assert "1772a4ea304041adb83f357b751341ff" not in item_list
+
+
+@pytest.mark.parametrize("chosen_env", ["m_adam_heating"], indirect=True)
+@pytest.mark.parametrize("cooling_present", [False], indirect=True)
+async def test_delete_removed_device(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_smile_adam_heat_cool: MagicMock,
+    device_registry: dr.DeviceRegistry,
+    init_integration: MockConfigEntry,
+) -> None:
+    """Test device removal at integration init."""
+    data = mock_smile_adam_heat_cool.async_update.return_value
+    mock_config_entry.add_to_hass(hass)
+    assert await async_setup_component(hass, DOMAIN, {})
+    await hass.async_block_till_done()
+
+    item_list: list[str] = []
+    for device_entry in device_registry.devices.values():
+        item_list.extend(x[1] for x in device_entry.identifiers)
+    assert "14df5c4dc8cb4ba69f9d1ac0eaf7c5c6" in item_list
+
+    data.pop("14df5c4dc8cb4ba69f9d1ac0eaf7c5c6")
+    with patch(HA_PLUGWISE_SMILE_ASYNC_UPDATE, return_value=data):
+        await hass.config_entries.async_reload(init_integration.entry_id)
+        await hass.async_block_till_done()
+
+    item_list = []
+    for device_entry in device_registry.devices.values():
+        item_list.extend(x[1] for x in device_entry.identifiers)
+    assert "14df5c4dc8cb4ba69f9d1ac0eaf7c5c6" not in item_list
+
+
+#### pw-beta only ####
+@pytest.mark.parametrize("chosen_env", ["m_anna_heatpump_cooling"], indirect=True)
+@pytest.mark.parametrize("cooling_present", [True], indirect=True)
+async def test_entry_migration(
+    hass: HomeAssistant, mock_smile_anna: MagicMock
+) -> None:
+    """Test config entry 1.2 -> 1.1 migration."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_HOST: "127.0.0.1",
+            CONF_MAC: "AA:BB:CC:DD:EE:FF",
+            CONF_PASSWORD: "test-password",
+            CONF_PORT: 80,
+            CONF_TIMEOUT: 30,
+            CONF_USERNAME: "smile",
+        },
+        minor_version=2,
+        version=1,
+        unique_id="smile98765",
+    )
+
+    entry.runtime_data = MagicMock(api=mock_smile_anna)
+    entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert entry.version == 1
+    assert entry.minor_version == 1
+    assert entry.data.get(CONF_TIMEOUT) is None
+    assert entry.state is ConfigEntryState.LOADED
