Dynamically add Airzone entities (#121891)

* airzone: reload entry on new devices Signed-off-by: Álvaro Fernández Rojas <[email protected]> * Revert "airzone: reload entry on new devices" This reverts commit 3ecc084. * airzone: sensor: dynamically add new entities Signed-off-by: Álvaro Fernández Rojas <[email protected]> * tests: restore reverted airzone tests Signed-off-by: Álvaro Fernández Rojas <[email protected]> * airzone: sensor: code fixes Signed-off-by: Álvaro Fernández Rojas <[email protected]> * airzone: water_heater: dynamically add entities Signed-off-by: Álvaro Fernández Rojas <[email protected]> * airzone: binary_sensor: dynamically add entities Signed-off-by: Álvaro Fernández Rojas <[email protected]> * airzone: select: dynamically add entities Signed-off-by: Álvaro Fernández Rojas <[email protected]> * airzone: climate: dynamically add entities Signed-off-by: Álvaro Fernández Rojas <[email protected]> * tests: airzone: use freezer Signed-off-by: Álvaro Fernández Rojas <[email protected]> * airzone: call async_add_entities once Signed-off-by: Álvaro Fernández Rojas <[email protected]> * airzone: wrap async_add_listener on async_on_unload Signed-off-by: Álvaro Fernández Rojas <[email protected]> * airzone: reduce number of entity listeners Signed-off-by: Álvaro Fernández Rojas <[email protected]> --------- Signed-off-by: Álvaro Fernández Rojas <[email protected]>
home-assistant · Jul 13, 2024 · c044417 · c044417
1 parent 3111951
commit c044417
Show file tree

Hide file tree

Showing 7 changed files with 205 additions and 67 deletions.
diff --git a/homeassistant/components/airzone/binary_sensor.py b/homeassistant/components/airzone/binary_sensor.py
@@ -82,33 +82,54 @@ async def async_setup_entry(
     """Add Airzone binary sensors from a config_entry."""
     coordinator = entry.runtime_data
 
-    binary_sensors: list[AirzoneBinarySensor] = [
-        AirzoneSystemBinarySensor(
-            coordinator,
-            description,
-            entry,
-            system_id,
-            system_data,
-        )
-        for system_id, system_data in coordinator.data[AZD_SYSTEMS].items()
-        for description in SYSTEM_BINARY_SENSOR_TYPES
-        if description.key in system_data
-    ]
-
-    binary_sensors.extend(
-        AirzoneZoneBinarySensor(
-            coordinator,
-            description,
-            entry,
-            system_zone_id,
-            zone_data,
-        )
-        for system_zone_id, zone_data in coordinator.data[AZD_ZONES].items()
-        for description in ZONE_BINARY_SENSOR_TYPES
-        if description.key in zone_data
-    )
-
-    async_add_entities(binary_sensors)
+    added_systems: set[str] = set()
+    added_zones: set[str] = set()
+
+    def _async_entity_listener() -> None:
+        """Handle additions of binary sensors."""
+
+        entities: list[AirzoneBinarySensor] = []
+
+        systems_data = coordinator.data.get(AZD_SYSTEMS, {})
+        received_systems = set(systems_data)
+        new_systems = received_systems - added_systems
+        if new_systems:
+            entities.extend(
+                AirzoneSystemBinarySensor(
+                    coordinator,
+                    description,
+                    entry,
+                    system_zone_id,
+                    systems_data.get(system_zone_id),
+                )
+                for system_zone_id in new_systems
+                for description in SYSTEM_BINARY_SENSOR_TYPES
+                if description.key in systems_data.get(system_zone_id)
+            )
+            added_systems.update(new_systems)
+
+        zones_data = coordinator.data.get(AZD_ZONES, {})
+        received_zones = set(zones_data)
+        new_zones = received_zones - added_zones
+        if new_zones:
+            entities.extend(
+                AirzoneZoneBinarySensor(
+                    coordinator,
+                    description,
+                    entry,
+                    system_zone_id,
+                    zones_data.get(system_zone_id),
+                )
+                for system_zone_id in new_zones
+                for description in ZONE_BINARY_SENSOR_TYPES
+                if description.key in zones_data.get(system_zone_id)
+            )
+            added_zones.update(new_zones)
+
+        async_add_entities(entities)
+
+    entry.async_on_unload(coordinator.async_add_listener(_async_entity_listener))
+    _async_entity_listener()
 
 
 class AirzoneBinarySensor(AirzoneEntity, BinarySensorEntity):

diff --git a/homeassistant/components/airzone/climate.py b/homeassistant/components/airzone/climate.py
@@ -102,17 +102,31 @@ async def async_setup_entry(
     entry: AirzoneConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
-    """Add Airzone sensors from a config_entry."""
+    """Add Airzone climate from a config_entry."""
     coordinator = entry.runtime_data
-    async_add_entities(
-        AirzoneClimate(
-            coordinator,
-            entry,
-            system_zone_id,
-            zone_data,
-        )
-        for system_zone_id, zone_data in coordinator.data[AZD_ZONES].items()
-    )
+
+    added_zones: set[str] = set()
+
+    def _async_entity_listener() -> None:
+        """Handle additions of climate."""
+
+        zones_data = coordinator.data.get(AZD_ZONES, {})
+        received_zones = set(zones_data)
+        new_zones = received_zones - added_zones
+        if new_zones:
+            async_add_entities(
+                AirzoneClimate(
+                    coordinator,
+                    entry,
+                    system_zone_id,
+                    zones_data.get(system_zone_id),
+                )
+                for system_zone_id in new_zones
+            )
+            added_zones.update(new_zones)
+
+    entry.async_on_unload(coordinator.async_add_listener(_async_entity_listener))
+    _async_entity_listener()
 
 
 class AirzoneClimate(AirzoneZoneEntity, ClimateEntity):

diff --git a/homeassistant/components/airzone/select.py b/homeassistant/components/airzone/select.py
@@ -83,21 +83,34 @@ async def async_setup_entry(
     entry: AirzoneConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
-    """Add Airzone sensors from a config_entry."""
+    """Add Airzone select from a config_entry."""
     coordinator = entry.runtime_data
 
-    async_add_entities(
-        AirzoneZoneSelect(
-            coordinator,
-            description,
-            entry,
-            system_zone_id,
-            zone_data,
-        )
-        for description in ZONE_SELECT_TYPES
-        for system_zone_id, zone_data in coordinator.data[AZD_ZONES].items()
-        if description.key in zone_data
-    )
+    added_zones: set[str] = set()
+
+    def _async_entity_listener() -> None:
+        """Handle additions of select."""
+
+        zones_data = coordinator.data.get(AZD_ZONES, {})
+        received_zones = set(zones_data)
+        new_zones = received_zones - added_zones
+        if new_zones:
+            async_add_entities(
+                AirzoneZoneSelect(
+                    coordinator,
+                    description,
+                    entry,
+                    system_zone_id,
+                    zones_data.get(system_zone_id),
+                )
+                for system_zone_id in new_zones
+                for description in ZONE_SELECT_TYPES
+                if description.key in zones_data.get(system_zone_id)
+            )
+            added_zones.update(new_zones)
+
+    entry.async_on_unload(coordinator.async_add_listener(_async_entity_listener))
+    _async_entity_listener()
 
 
 class AirzoneBaseSelect(AirzoneEntity, SelectEntity):

diff --git a/homeassistant/components/airzone/sensor.py b/homeassistant/components/airzone/sensor.py
@@ -85,21 +85,37 @@ async def async_setup_entry(
     """Add Airzone sensors from a config_entry."""
     coordinator = entry.runtime_data
 
-    sensors: list[AirzoneSensor] = [
-        AirzoneZoneSensor(
-            coordinator,
-            description,
-            entry,
-            system_zone_id,
-            zone_data,
-        )
-        for system_zone_id, zone_data in coordinator.data[AZD_ZONES].items()
-        for description in ZONE_SENSOR_TYPES
-        if description.key in zone_data
-    ]
+    added_zones: set[str] = set()
+
+    def _async_entity_listener() -> None:
+        """Handle additions of sensors."""
+
+        entities: list[AirzoneSensor] = []
+
+        zones_data = coordinator.data.get(AZD_ZONES, {})
+        received_zones = set(zones_data)
+        new_zones = received_zones - added_zones
+        if new_zones:
+            entities.extend(
+                AirzoneZoneSensor(
+                    coordinator,
+                    description,
+                    entry,
+                    system_zone_id,
+                    zones_data.get(system_zone_id),
+                )
+                for system_zone_id in new_zones
+                for description in ZONE_SENSOR_TYPES
+                if description.key in zones_data.get(system_zone_id)
+            )
+            added_zones.update(new_zones)
+
+        async_add_entities(entities)
+
+    entities: list[AirzoneSensor] = []
 
     if AZD_HOT_WATER in coordinator.data:
-        sensors.extend(
+        entities.extend(
             AirzoneHotWaterSensor(
                 coordinator,
                 description,
@@ -110,7 +126,7 @@ async def async_setup_entry(
         )
 
     if AZD_WEBSERVER in coordinator.data:
-        sensors.extend(
+        entities.extend(
             AirzoneWebServerSensor(
                 coordinator,
                 description,
@@ -120,7 +136,10 @@ async def async_setup_entry(
             if description.key in coordinator.data[AZD_WEBSERVER]
         )
 
-    async_add_entities(sensors)
+    async_add_entities(entities)
+
+    entry.async_on_unload(coordinator.async_add_listener(_async_entity_listener))
+    _async_entity_listener()
 
 
 class AirzoneSensor(AirzoneEntity, SensorEntity):

diff --git a/homeassistant/components/airzone/water_heater.py b/homeassistant/components/airzone/water_heater.py
@@ -61,7 +61,7 @@ async def async_setup_entry(
     entry: AirzoneConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
-    """Add Airzone sensors from a config_entry."""
+    """Add Airzone Water Heater from a config_entry."""
     coordinator = entry.runtime_data
     if AZD_HOT_WATER in coordinator.data:
         async_add_entities([AirzoneWaterHeater(coordinator, entry)])

diff --git a/tests/components/airzone/test_coordinator.py b/tests/components/airzone/test_coordinator.py
@@ -8,14 +8,15 @@
     InvalidMethod,
     SystemOutOfRange,
 )
+from freezegun.api import FrozenDateTimeFactory
 
 from homeassistant.components.airzone.const import DOMAIN
 from homeassistant.components.airzone.coordinator import SCAN_INTERVAL
 from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 from homeassistant.util.dt import utcnow
 
-from .util import CONFIG, HVAC_MOCK, HVAC_VERSION_MOCK
+from .util import CONFIG, HVAC_MOCK, HVAC_MOCK_NEW_ZONES, HVAC_VERSION_MOCK
 
 from tests.common import MockConfigEntry, async_fire_time_changed
 
@@ -64,3 +65,62 @@ async def test_coordinator_client_connector_error(hass: HomeAssistant) -> None:
 
         state = hass.states.get("sensor.despacho_temperature")
         assert state.state == STATE_UNAVAILABLE
+
+
+async def test_coordinator_new_devices(
+    hass: HomeAssistant,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test new devices on coordinator update."""
+
+    config_entry = MockConfigEntry(
+        data=CONFIG,
+        domain=DOMAIN,
+        unique_id="airzone_unique_id",
+    )
+    config_entry.add_to_hass(hass)
+
+    with (
+        patch(
+            "homeassistant.components.airzone.AirzoneLocalApi.get_dhw",
+            side_effect=HotWaterNotAvailable,
+        ),
+        patch(
+            "homeassistant.components.airzone.AirzoneLocalApi.get_hvac",
+            return_value=HVAC_MOCK_NEW_ZONES,
+        ) as mock_hvac,
+        patch(
+            "homeassistant.components.airzone.AirzoneLocalApi.get_hvac_systems",
+            side_effect=SystemOutOfRange,
+        ),
+        patch(
+            "homeassistant.components.airzone.AirzoneLocalApi.get_version",
+            return_value=HVAC_VERSION_MOCK,
+        ),
+        patch(
+            "homeassistant.components.airzone.AirzoneLocalApi.get_webserver",
+            side_effect=InvalidMethod,
+        ),
+    ):
+        await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+        mock_hvac.assert_called_once()
+        mock_hvac.reset_mock()
+
+        state = hass.states.get("sensor.salon_temperature")
+        assert state.state == "19.6"
+
+        state = hass.states.get("sensor.dorm_ppal_temperature")
+        assert state is None
+
+        mock_hvac.return_value = HVAC_MOCK
+        freezer.tick(SCAN_INTERVAL)
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
+        mock_hvac.assert_called_once()
+
+        state = hass.states.get("sensor.salon_temperature")
+        assert state.state == "19.6"
+
+        state = hass.states.get("sensor.dorm_ppal_temperature")
+        assert state.state == "21.1"
diff --git a/tests/components/airzone/util.py b/tests/components/airzone/util.py
@@ -1,5 +1,6 @@
 """Tests for the Airzone integration."""
 
+from copy import deepcopy
 from unittest.mock import patch
 
 from aioairzone.const import (
@@ -274,6 +275,16 @@
     ]
 }
 
+HVAC_MOCK_NEW_ZONES = {
+    API_SYSTEMS: [
+        {
+            API_DATA: [
+                deepcopy(HVAC_MOCK[API_SYSTEMS][0][API_DATA][0]),
+            ]
+        }
+    ]
+}
+
 HVAC_DHW_MOCK = {
     API_DATA: {
         API_SYSTEM_ID: 0,