From 3ecc0844e42a9c088a0acb1fea085cb1a84f6c5a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81lvaro=20Fern=C3=A1ndez=20Rojas?= <noltari@gmail.com>
Date: Sat, 13 Jul 2024 09:14:43 +0200
Subject: [PATCH] airzone: reload entry on new devices
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Álvaro Fernández Rojas <noltari@gmail.com>
---
 .../components/airzone/coordinator.py         | 14 ++++-
 tests/components/airzone/test_coordinator.py  | 57 ++++++++++++++++++-
 tests/components/airzone/util.py              | 11 ++++
 3 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/airzone/coordinator.py b/homeassistant/components/airzone/coordinator.py
index 8ec2cbe07caaca..5df6a95e8eeedf 100644
--- a/homeassistant/components/airzone/coordinator.py
+++ b/homeassistant/components/airzone/coordinator.py
@@ -7,6 +7,7 @@
 import logging
 from typing import Any
 
+from aioairzone.const import AZD_NEW_SYSTEMS, AZD_NEW_ZONES
 from aioairzone.exceptions import AirzoneError
 from aioairzone.localapi import AirzoneLocalApi
 
@@ -34,6 +35,15 @@ def __init__(self, hass: HomeAssistant, airzone: AirzoneLocalApi) -> None:
             update_interval=SCAN_INTERVAL,
         )
 
+    async def _async_check_new_devices(self, data: dict[str, Any]) -> None:
+        """Reload entry if there are new devices."""
+        if (
+            len(data.get(AZD_NEW_SYSTEMS, [])) > 0
+            or len(data.get(AZD_NEW_ZONES, [])) > 0
+        ):
+            assert self.config_entry
+            self.hass.config_entries.async_schedule_reload(self.config_entry.entry_id)
+
     async def _async_update_data(self) -> dict[str, Any]:
         """Update data via library."""
         async with timeout(AIOAIRZONE_DEVICE_TIMEOUT_SEC):
@@ -41,4 +51,6 @@ async def _async_update_data(self) -> dict[str, Any]:
                 await self.airzone.update()
             except AirzoneError as error:
                 raise UpdateFailed(error) from error
-            return self.airzone.data()
+            data = self.airzone.data()
+        await self._async_check_new_devices(data)
+        return data
diff --git a/tests/components/airzone/test_coordinator.py b/tests/components/airzone/test_coordinator.py
index 06c77bebb81747..f42b2e88cadaa0 100644
--- a/tests/components/airzone/test_coordinator.py
+++ b/tests/components/airzone/test_coordinator.py
@@ -15,7 +15,7 @@
 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 +64,58 @@ 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) -> 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
+        async_fire_time_changed(hass, utcnow() + SCAN_INTERVAL)
+        await hass.async_block_till_done()
+        mock_hvac.assert_called()
+
+        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
index 6e3e0eccc8f968..2cdb7a9c6f9c1f 100644
--- 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,