diff --git a/homeassistant/components/husqvarna_automower/coordinator.py b/homeassistant/components/husqvarna_automower/coordinator.py index 70af5219d04..342f6892b2e 100644 --- a/homeassistant/components/husqvarna_automower/coordinator.py +++ b/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]] = {} - - 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) + self.async_add_listener(self._on_data_update) 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 } diff --git a/tests/components/husqvarna_automower/test_init.py b/tests/components/husqvarna_automower/test_init.py index 9a45b2ad42d..f54250a3336 100644 --- a/tests/components/husqvarna_automower/test_init.py +++ b/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") - values[TEST_MOWER_ID].work_area_dict.update({1: "new work area"}) - values[TEST_MOWER_ID].work_areas.update( + await hass.async_block_till_done() + + 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 - freezer.tick(SCAN_INTERVAL) - async_fire_time_changed(hass) + mock_automower_client.get_status.return_value = poll_values + + 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") - del values[TEST_MOWER_ID].work_area_dict[1] - del values[TEST_MOWER_ID].work_areas[1] - values[TEST_MOWER_ID].work_area_names.remove("Front lawn") - del values[TEST_MOWER_ID].work_area_dict[123456] - del values[TEST_MOWER_ID].work_areas[123456] - del values[TEST_MOWER_ID].calendar.tasks[:2] - values[TEST_MOWER_ID].mower.work_area_id = 654321 - mock_automower_client.get_status.return_value = values + poll_values[TEST_MOWER_ID].work_area_names.remove("new work area") + del poll_values[TEST_MOWER_ID].work_area_dict[1] + del poll_values[TEST_MOWER_ID].work_areas[1] + poll_values[TEST_MOWER_ID].work_area_names.remove("Front lawn") + del poll_values[TEST_MOWER_ID].work_area_dict[123456] + del poll_values[TEST_MOWER_ID].work_areas[123456] + del poll_values[TEST_MOWER_ID].calendar.tasks[:2] + poll_values[TEST_MOWER_ID].mower.work_area_id = 654321 + mock_automower_client.get_status.return_value = poll_values freezer.tick(SCAN_INTERVAL) async_fire_time_changed(hass) await hass.async_block_till_done()