Reduce polling in Husqvarna Automower (#149255)

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

homeassistant/components/husqvarna_automower/coordinator.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import asyncio
 from collections.abc import Callable
-from datetime import timedelta
+from datetime import datetime, timedelta
 import logging
 from typing import override
 
@@ -14,7 +14,7 @@ from aioautomower.exceptions import (
     HusqvarnaTimeoutError,
     HusqvarnaWSServerHandshakeError,
 )
-from aioautomower.model import MowerDictionary
+from aioautomower.model import MowerDictionary, MowerStates
 from aioautomower.session import AutomowerSession
 
 from homeassistant.config_entries import ConfigEntry
@@ -29,7 +29,9 @@ _LOGGER = logging.getLogger(__name__)
 MAX_WS_RECONNECT_TIME = 600
 SCAN_INTERVAL = timedelta(minutes=8)
 DEFAULT_RECONNECT_TIME = 2  # Define a default reconnect time
-
+PONG_TIMEOUT = timedelta(seconds=90)
+PING_INTERVAL = timedelta(seconds=10)
+PING_TIMEOUT = timedelta(seconds=5)
 type AutomowerConfigEntry = ConfigEntry[AutomowerDataUpdateCoordinator]
 
 
@@ -58,6 +60,9 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[MowerDictionary]):
         self.new_devices_callbacks: list[Callable[[set[str]], None]] = []
         self.new_zones_callbacks: list[Callable[[str, set[str]], None]] = []
         self.new_areas_callbacks: list[Callable[[str, set[int]], None]] = []
+        self.pong: datetime | None = None
+        self.websocket_alive: bool = False
+        self._watchdog_task: asyncio.Task | None = None
 
     @override
     @callback
@@ -71,6 +76,18 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[MowerDictionary]):
             await self.api.connect()
             self.api.register_data_callback(self.handle_websocket_updates)
             self.ws_connected = True
+
+            def start_watchdog() -> None:
+                if self._watchdog_task is not None and not self._watchdog_task.done():
+                    _LOGGER.debug("Cancelling previous watchdog task")
+                    self._watchdog_task.cancel()
+                self._watchdog_task = self.config_entry.async_create_background_task(
+                    self.hass,
+                    self._pong_watchdog(),
+                    "websocket_watchdog",
+                )
+
+            self.api.register_ws_ready_callback(start_watchdog)
         try:
             data = await self.api.get_status()
         except ApiError as err:
@@ -93,6 +110,19 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[MowerDictionary]):
                 mower_data.capabilities.work_areas for mower_data in self.data.values()
             ):
                 self._async_add_remove_work_areas()
+            if (
+                not self._should_poll()
+                and self.update_interval is not None
+                and self.websocket_alive
+            ):
+                _LOGGER.debug("All mowers inactive and websocket alive: stop polling")
+                self.update_interval = None
+            if self.update_interval is None and self._should_poll():
+                _LOGGER.debug(
+                    "Polling re-enabled via WebSocket: at least one mower active"
+                )
+                self.update_interval = SCAN_INTERVAL
+                self.hass.async_create_task(self.async_request_refresh())
 
     @callback
     def handle_websocket_updates(self, ws_data: MowerDictionary) -> None:
@@ -161,6 +191,30 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[MowerDictionary]):
                 "reconnect_task",
             )
 
+    def _should_poll(self) -> bool:
+        """Return True if at least one mower is connected and at least one is not OFF."""
+        return any(mower.metadata.connected for mower in self.data.values()) and any(
+            mower.mower.state != MowerStates.OFF for mower in self.data.values()
+        )
+
+    async def _pong_watchdog(self) -> None:
+        _LOGGER.debug("Watchdog started")
+        try:
+            while True:
+                _LOGGER.debug("Sending ping")
+                self.websocket_alive = await self.api.send_empty_message()
+                _LOGGER.debug("Ping result: %s", self.websocket_alive)
+
+                await asyncio.sleep(60)
+                _LOGGER.debug("Websocket alive %s", self.websocket_alive)
+                if not self.websocket_alive:
+                    _LOGGER.debug("No pong received → restart polling")
+                    if self.update_interval is None:
+                        self.update_interval = SCAN_INTERVAL
+                        await self.async_request_refresh()
+        except asyncio.CancelledError:
+            _LOGGER.debug("Watchdog cancelled")
+
     def _async_add_remove_devices(self) -> None:
         """Add new devices and remove orphaned devices from the registry."""
         current_devices = set(self.data)

tests/components/husqvarna_automower/test_init.py

@@ -14,7 +14,7 @@ from aioautomower.exceptions import (
     HusqvarnaTimeoutError,
     HusqvarnaWSServerHandshakeError,
 )
-from aioautomower.model import Calendar, MowerAttributes, WorkArea
+from aioautomower.model import Calendar, MowerAttributes, MowerStates, WorkArea
 from freezegun.api import FrozenDateTimeFactory
 import pytest
 from syrupy.assertion import SnapshotAssertion
@@ -484,3 +484,212 @@ async def test_add_and_remove_work_area(
         - ADDITIONAL_NUMBER_ENTITIES
         - ADDITIONAL_SENSOR_ENTITIES
     )
+
+
+@pytest.mark.parametrize(
+    ("mower1_connected", "mower1_state", "mower2_connected", "mower2_state"),
+    [
+        (True, MowerStates.OFF, False, MowerStates.OFF),  # False
+        (False, MowerStates.PAUSED, False, MowerStates.OFF),  # False
+        (False, MowerStates.OFF, True, MowerStates.OFF),  # False
+        (False, MowerStates.OFF, False, MowerStates.PAUSED),  # False
+        (True, MowerStates.OFF, True, MowerStates.OFF),  # False
+        (False, MowerStates.OFF, False, MowerStates.OFF),  # False
+    ],
+)
+async def test_dynamic_polling(
+    hass: HomeAssistant,
+    mock_automower_client,
+    mock_config_entry,
+    freezer: FrozenDateTimeFactory,
+    values: dict[str, MowerAttributes],
+    mower1_connected: bool,
+    mower1_state: MowerStates,
+    mower2_connected: bool,
+    mower2_state: MowerStates,
+) -> None:
+    """Test that the ws_ready_callback triggers an attempt to start the Watchdog task.
+
+    and that the pong callback stops polling when all mowers are inactive.
+    """
+    websocket_values = deepcopy(values)
+    poll_values = deepcopy(values)
+    callback_holder: dict[str, Callable] = {}
+
+    @callback
+    def fake_register_websocket_response(
+        cb: Callable[[dict[str, MowerAttributes]], None],
+    ) -> None:
+        callback_holder["data_cb"] = cb
+
+    mock_automower_client.register_data_callback.side_effect = (
+        fake_register_websocket_response
+    )
+
+    @callback
+    def fake_register_ws_ready_callback(cb: Callable[[], None]) -> None:
+        callback_holder["ws_ready_cb"] = cb
+
+    mock_automower_client.register_ws_ready_callback.side_effect = (
+        fake_register_ws_ready_callback
+    )
+
+    await setup_integration(hass, mock_config_entry)
+
+    assert "ws_ready_cb" in callback_holder, "ws_ready_callback was not registered"
+    callback_holder["ws_ready_cb"]()
+
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 1
+
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 2
+
+    # websocket is still active, but mowers are inactive -> no polling required
+    poll_values[TEST_MOWER_ID].metadata.connected = mower1_connected
+    poll_values[TEST_MOWER_ID].mower.state = mower1_state
+    poll_values["1234"].metadata.connected = mower2_connected
+    poll_values["1234"].mower.state = mower2_state
+
+    mock_automower_client.get_status.return_value = poll_values
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 3
+
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 4
+
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 4
+
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 4
+
+    # websocket is still active, and mowers are active -> polling required
+    mock_automower_client.get_status.reset_mock()
+    assert mock_automower_client.get_status.call_count == 0
+    poll_values[TEST_MOWER_ID].metadata.connected = True
+    poll_values[TEST_MOWER_ID].mower.state = MowerStates.PAUSED
+    poll_values["1234"].metadata.connected = False
+    poll_values["1234"].mower.state = MowerStates.OFF
+    websocket_values = deepcopy(poll_values)
+    callback_holder["data_cb"](websocket_values)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 1
+
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 2
+
+
+@pytest.mark.parametrize(
+    ("mower1_connected", "mower1_state", "mower2_connected", "mower2_state"),
+    [
+        (True, MowerStates.OFF, False, MowerStates.OFF),  # False
+        (False, MowerStates.PAUSED, False, MowerStates.OFF),  # False
+        (False, MowerStates.OFF, True, MowerStates.OFF),  # False
+        (False, MowerStates.OFF, False, MowerStates.PAUSED),  # False
+        (True, MowerStates.OFF, True, MowerStates.OFF),  # False
+        (False, MowerStates.OFF, False, MowerStates.OFF),  # False
+    ],
+)
+async def test_websocket_watchdog(
+    hass: HomeAssistant,
+    mock_automower_client,
+    mock_config_entry,
+    freezer: FrozenDateTimeFactory,
+    entity_registry: er.EntityRegistry,
+    values: dict[str, MowerAttributes],
+    mower1_connected: bool,
+    mower1_state: MowerStates,
+    mower2_connected: bool,
+    mower2_state: MowerStates,
+) -> None:
+    """Test that the ws_ready_callback triggers an attempt to start the Watchdog task.
+
+    and that the pong callback stops polling when all mowers are inactive.
+    """
+    poll_values = deepcopy(values)
+    callback_holder: dict[str, Callable] = {}
+
+    @callback
+    def fake_register_websocket_response(
+        cb: Callable[[dict[str, MowerAttributes]], None],
+    ) -> None:
+        callback_holder["data_cb"] = cb
+
+    mock_automower_client.register_data_callback.side_effect = (
+        fake_register_websocket_response
+    )
+
+    @callback
+    def fake_register_ws_ready_callback(cb: Callable[[], None]) -> None:
+        callback_holder["ws_ready_cb"] = cb
+
+    mock_automower_client.register_ws_ready_callback.side_effect = (
+        fake_register_ws_ready_callback
+    )
+
+    await setup_integration(hass, mock_config_entry)
+
+    assert "ws_ready_cb" in callback_holder, "ws_ready_callback was not registered"
+    callback_holder["ws_ready_cb"]()
+
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 1
+
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 2
+
+    # websocket is still active, but mowers are inactive -> no polling required
+    poll_values[TEST_MOWER_ID].metadata.connected = mower1_connected
+    poll_values[TEST_MOWER_ID].mower.state = mower1_state
+    poll_values["1234"].metadata.connected = mower2_connected
+    poll_values["1234"].mower.state = mower2_state
+
+    mock_automower_client.get_status.return_value = poll_values
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 3
+
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 4
+
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 4
+
+    # Simulate Pong loss and reset mock -> polling required
+    mock_automower_client.send_empty_message.return_value = False
+    mock_automower_client.get_status.reset_mock()
+
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 0
+
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 1
+
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.get_status.call_count == 2