Fix adding a work area in Husqvarna Automower (#148358)

homeassistant/components/husqvarna_automower/coordinator.py

@@ -60,15 +60,7 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[MowerDictionary]):
         self._devices_last_update: set[str] = set()
         self._zones_last_update: dict[str, set[str]] = {}
         self._areas_last_update: dict[str, set[int]] = {}
-
+        self.async_add_listener(self._on_data_update)
-    def _async_add_remove_devices_and_entities(self, data: MowerDictionary) -> None:
-        """Add/remove devices and dynamic entities, when amount of devices changed."""
-        self._async_add_remove_devices(data)
-        for mower_id in data:
-            if data[mower_id].capabilities.stay_out_zones:
-                self._async_add_remove_stay_out_zones(data)
-            if data[mower_id].capabilities.work_areas:
-                self._async_add_remove_work_areas(data)
 
     async def _async_update_data(self) -> MowerDictionary:
         """Subscribe for websocket and poll data from the API."""
@@ -82,14 +74,38 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[MowerDictionary]):
             raise UpdateFailed(err) from err
         except AuthError as err:
             raise ConfigEntryAuthFailed(err) from err
-        self._async_add_remove_devices_and_entities(data)
         return data
 
+    @callback
+    def _on_data_update(self) -> None:
+        """Handle data updates and process dynamic entity management."""
+        if self.data is not None:
+            self._async_add_remove_devices()
+            for mower_id in self.data:
+                if self.data[mower_id].capabilities.stay_out_zones:
+                    self._async_add_remove_stay_out_zones()
+                if self.data[mower_id].capabilities.work_areas:
+                    self._async_add_remove_work_areas()
+
     @callback
     def handle_websocket_updates(self, ws_data: MowerDictionary) -> None:
         """Process websocket callbacks and write them to the DataUpdateCoordinator."""
+        self.hass.async_create_task(self._process_websocket_update(ws_data))
+
+    async def _process_websocket_update(self, ws_data: MowerDictionary) -> None:
+        """Handle incoming websocket update and update coordinator data."""
+        for data in ws_data.values():
+            existing_areas = data.work_areas or {}
+            for task in data.calendar.tasks:
+                work_area_id = task.work_area_id
+                if work_area_id is not None and work_area_id not in existing_areas:
+                    _LOGGER.debug(
+                        "New work area %s detected, refreshing data", work_area_id
+                    )
+                    await self.async_request_refresh()
+                    return
+
         self.async_set_updated_data(ws_data)
-        self._async_add_remove_devices_and_entities(ws_data)
 
     @callback
     def async_set_updated_data(self, data: MowerDictionary) -> None:
@@ -138,9 +154,9 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[MowerDictionary]):
                 "reconnect_task",
             )
 
-    def _async_add_remove_devices(self, data: MowerDictionary) -> None:
+    def _async_add_remove_devices(self) -> None:
         """Add new device, remove non-existing device."""
-        current_devices = set(data)
+        current_devices = set(self.data)
 
         # Skip update if no changes
         if current_devices == self._devices_last_update:
@@ -155,7 +171,6 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[MowerDictionary]):
         # Process new device
         new_devices = current_devices - self._devices_last_update
         if new_devices:
-            self.data = data
             _LOGGER.debug("New devices found: %s", ", ".join(map(str, new_devices)))
             self._add_new_devices(new_devices)
 
@@ -179,11 +194,11 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[MowerDictionary]):
         for mower_callback in self.new_devices_callbacks:
             mower_callback(new_devices)
 
-    def _async_add_remove_stay_out_zones(self, data: MowerDictionary) -> None:
+    def _async_add_remove_stay_out_zones(self) -> None:
         """Add new stay-out zones, remove non-existing stay-out zones."""
         current_zones = {
             mower_id: set(mower_data.stay_out_zones.zones)
-            for mower_id, mower_data in data.items()
+            for mower_id, mower_data in self.data.items()
             if mower_data.capabilities.stay_out_zones
             and mower_data.stay_out_zones is not None
         }
@@ -225,11 +240,11 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[MowerDictionary]):
 
         return current_zones
 
-    def _async_add_remove_work_areas(self, data: MowerDictionary) -> None:
+    def _async_add_remove_work_areas(self) -> None:
         """Add new work areas, remove non-existing work areas."""
         current_areas = {
             mower_id: set(mower_data.work_areas)
-            for mower_id, mower_data in data.items()
+            for mower_id, mower_data in self.data.items()
             if mower_data.capabilities.work_areas and mower_data.work_areas is not None
         }
 

tests/components/husqvarna_automower/test_init.py

@@ -3,7 +3,7 @@
 from asyncio import Event
 from collections.abc import Callable
 from copy import deepcopy
-from datetime import datetime, timedelta
+from datetime import datetime, time as dt_time, timedelta
 import http
 import time
 from unittest.mock import AsyncMock, patch
@@ -14,7 +14,7 @@ from aioautomower.exceptions import (
     HusqvarnaTimeoutError,
     HusqvarnaWSServerHandshakeError,
 )
-from aioautomower.model import MowerAttributes, WorkArea
+from aioautomower.model import Calendar, MowerAttributes, WorkArea
 from freezegun.api import FrozenDateTimeFactory
 import pytest
 from syrupy.assertion import SnapshotAssertion
@@ -384,14 +384,45 @@ async def test_add_and_remove_work_area(
     values: dict[str, MowerAttributes],
 ) -> None:
     """Test adding a work area in runtime."""
+    websocket_values = deepcopy(values)
+    callback_holder: dict[str, Callable] = {}
+
+    @callback
+    def fake_register_websocket_response(
+        cb: Callable[[dict[str, MowerAttributes]], None],
+    ) -> None:
+        callback_holder["cb"] = cb
+
+    mock_automower_client.register_data_callback.side_effect = (
+        fake_register_websocket_response
+    )
     await setup_integration(hass, mock_config_entry)
     entry = hass.config_entries.async_entries(DOMAIN)[0]
     current_entites_start = len(
         er.async_entries_for_config_entry(entity_registry, entry.entry_id)
     )
-    values[TEST_MOWER_ID].work_area_names.append("new work area")
+    await hass.async_block_till_done()
-    values[TEST_MOWER_ID].work_area_dict.update({1: "new work area"})
+
-    values[TEST_MOWER_ID].work_areas.update(
+    assert mock_automower_client.register_data_callback.called
+    assert "cb" in callback_holder
+
+    new_task = Calendar(
+        start=dt_time(hour=11),
+        duration=timedelta(60),
+        monday=True,
+        tuesday=True,
+        wednesday=True,
+        thursday=True,
+        friday=True,
+        saturday=True,
+        sunday=True,
+        work_area_id=1,
+    )
+    websocket_values[TEST_MOWER_ID].calendar.tasks.append(new_task)
+    poll_values = deepcopy(websocket_values)
+    poll_values[TEST_MOWER_ID].work_area_names.append("new work area")
+    poll_values[TEST_MOWER_ID].work_area_dict.update({1: "new work area"})
+    poll_values[TEST_MOWER_ID].work_areas.update(
         {
             1: WorkArea(
                 name="new work area",
@@ -404,10 +435,15 @@ async def test_add_and_remove_work_area(
             )
         }
     )
-    mock_automower_client.get_status.return_value = values
+    mock_automower_client.get_status.return_value = poll_values
-    freezer.tick(SCAN_INTERVAL)
+
-    async_fire_time_changed(hass)
+    callback_holder["cb"](websocket_values)
     await hass.async_block_till_done()
+    assert mock_automower_client.get_status.called
+
+    state = hass.states.get("sensor.test_mower_1_new_work_area_progress")
+    assert state is not None
+    assert state.state == "12"
     current_entites_after_addition = len(
         er.async_entries_for_config_entry(entity_registry, entry.entry_id)
     )
@@ -419,15 +455,15 @@ async def test_add_and_remove_work_area(
         + ADDITIONAL_SWITCH_ENTITIES
     )
 
-    values[TEST_MOWER_ID].work_area_names.remove("new work area")
+    poll_values[TEST_MOWER_ID].work_area_names.remove("new work area")
-    del values[TEST_MOWER_ID].work_area_dict[1]
+    del poll_values[TEST_MOWER_ID].work_area_dict[1]
-    del values[TEST_MOWER_ID].work_areas[1]
+    del poll_values[TEST_MOWER_ID].work_areas[1]
-    values[TEST_MOWER_ID].work_area_names.remove("Front lawn")
+    poll_values[TEST_MOWER_ID].work_area_names.remove("Front lawn")
-    del values[TEST_MOWER_ID].work_area_dict[123456]
+    del poll_values[TEST_MOWER_ID].work_area_dict[123456]
-    del values[TEST_MOWER_ID].work_areas[123456]
+    del poll_values[TEST_MOWER_ID].work_areas[123456]
-    del values[TEST_MOWER_ID].calendar.tasks[:2]
+    del poll_values[TEST_MOWER_ID].calendar.tasks[:2]
-    values[TEST_MOWER_ID].mower.work_area_id = 654321
+    poll_values[TEST_MOWER_ID].mower.work_area_id = 654321
-    mock_automower_client.get_status.return_value = values
+    mock_automower_client.get_status.return_value = poll_values
     freezer.tick(SCAN_INTERVAL)
     async_fire_time_changed(hass)
     await hass.async_block_till_done()