Improve Plugwise coordinator code (#158983)

homeassistant/components/plugwise/coordinator.py

@@ -7,6 +7,7 @@ from plugwise import GwEntityData, Smile
 from plugwise.exceptions import (
     ConnectionFailedError,
     InvalidAuthentication,
+    InvalidSetupError,
     InvalidXMLError,
     PlugwiseError,
     ResponseError,
@@ -31,6 +32,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
 
@@ -59,14 +63,31 @@ class PlugwiseDataUpdateCoordinator(DataUpdateCoordinator[dict[str, GwEntityData
             port=self.config_entry.data.get(CONF_PORT, DEFAULT_PORT),
             websession=async_get_clientsession(hass, verify_ssl=False),
         )
-        self._current_devices: set[str] = set()
+        self._current_devices = set()
-        self.new_devices: set[str] = set()
+        self._stored_devices = set()
+        self.new_devices = set()
 
     async def _connect(self) -> None:
-        """Connect to the Plugwise Smile."""
+        """Connect to the Plugwise Smile.
+
+        A Version object is received when the connection succeeds.
+        """
         version = await self.api.connect()
         self._connected = isinstance(version, Version)
 
+    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."""
         try:
@@ -83,10 +104,15 @@ class PlugwiseDataUpdateCoordinator(DataUpdateCoordinator[dict[str, GwEntityData
                 translation_domain=DOMAIN,
                 translation_key="authentication_failed",
             ) from err
+        except InvalidSetupError as err:
+            raise ConfigEntryError(
+                translation_domain=DOMAIN,
+                translation_key="invalid_setup",
+            ) from err
         except (InvalidXMLError, ResponseError) as err:
             raise UpdateFailed(
                 translation_domain=DOMAIN,
-                translation_key="invalid_xml_data",
+                translation_key="response_error",
             ) from err
         except PlugwiseError as err:
             raise UpdateFailed(
@@ -104,12 +130,16 @@ class PlugwiseDataUpdateCoordinator(DataUpdateCoordinator[dict[str, GwEntityData
 
     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
+        set_of_data = set(data)
-        self.new_devices = set(data) - self._current_devices
+        # Check for new or removed devices,
-        removed_devices = self._current_devices - set(data)
+        # 'new_devices' contains all devices present in 'data' at init ('self._current_devices' is empty)
-        self._current_devices = set(data)
+        # this is required for the proper initialization of all the present platform entities.
-
+        self.new_devices = set_of_data - self._current_devices
-        if removed_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
             self._async_remove_devices(data)
 
     def _async_remove_devices(self, data: dict[str, GwEntityData]) -> None:
@@ -118,17 +148,17 @@ class PlugwiseDataUpdateCoordinator(DataUpdateCoordinator[dict[str, GwEntityData
         device_list = dr.async_entries_for_config_entry(
             device_reg, self.config_entry.entry_id
         )
+
         # First find the Plugwise via_device
         gateway_device = device_reg.async_get_device({(DOMAIN, self.api.gateway_id)})
         assert gateway_device is not None
         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:
-                if identifier[0] == DOMAIN:
                 if (
-                        device_entry.via_device_id == via_device_id
+                    identifier[0] == DOMAIN
+                    and device_entry.via_device_id == via_device_id
                     and identifier[1] not in data
                 ):
                     device_reg.async_update_device(
@@ -136,7 +166,7 @@ class PlugwiseDataUpdateCoordinator(DataUpdateCoordinator[dict[str, GwEntityData
                         remove_config_entry_id=self.config_entry.entry_id,
                     )
                     LOGGER.debug(
-                            "Removed %s device %s %s from device_registry",
+                        "Removed %s device/zone %s %s from device_registry",
                         DOMAIN,
                         device_entry.model,
                         identifier[1],

homeassistant/components/plugwise/strings.json

@@ -319,7 +319,10 @@
     "failed_to_connect": {
       "message": "[%key:common::config_flow::error::cannot_connect%]"
     },
-    "invalid_xml_data": {
+    "invalid_setup": {
+      "message": "Add your Adam instead of your Anna, see the documentation"
+    },
+    "response_error": {
       "message": "[%key:component::plugwise::config::error::response_error%]"
     },
     "set_schedule_first": {

tests/components/plugwise/test_init.py

@@ -7,6 +7,7 @@ from freezegun.api import FrozenDateTimeFactory
 from plugwise.exceptions import (
     ConnectionFailedError,
     InvalidAuthentication,
+    InvalidSetupError,
     InvalidXMLError,
     PlugwiseError,
     ResponseError,
@@ -89,6 +90,7 @@ async def test_load_unload_config_entry(
     [
         (ConnectionFailedError, ConfigEntryState.SETUP_RETRY),
         (InvalidAuthentication, ConfigEntryState.SETUP_ERROR),
+        (InvalidSetupError, ConfigEntryState.SETUP_ERROR),
         (InvalidXMLError, ConfigEntryState.SETUP_RETRY),
         (PlugwiseError, ConfigEntryState.SETUP_RETRY),
         (ResponseError, ConfigEntryState.SETUP_RETRY),
@@ -169,7 +171,7 @@ async def test_migrate_unique_id_temperature(
     """Test migration of unique_id."""
     mock_config_entry.add_to_hass(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,
     )
@@ -334,3 +336,31 @@ async def test_update_device(
         for device_entry in list(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
+
+    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