Avoid errors when there is no internet connection in Husqvarna Automo…

…wer (#111101)

* Avoid errors when no internet connection

* Add error

* Create task in HA

* change from matter to automower

* tests

* Update homeassistant/components/husqvarna_automower/coordinator.py

Co-authored-by: Martin Hjelmare <[email protected]>

* address review

* Make websocket optional

* fix aioautomower version

* Fix tests

* Use stored websocket

* reset reconnect time after sucessful connection

* Typo

* Remove comment

* Add test

* Address review

---------

Co-authored-by: Martin Hjelmare <[email protected]>

homeassistant/components/husqvarna_automower/__init__.py

@@ -6,7 +6,7 @@
 from aiohttp import ClientError
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import EVENT_HOMEASSISTANT_STOP, Platform
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import aiohttp_client, config_entry_oauth2_flow
@@ -17,7 +17,6 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-
 PLATFORMS: list[Platform] = [Platform.LAWN_MOWER, Platform.SENSOR, Platform.SWITCH]
 
 
@@ -38,22 +37,22 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         await api_api.async_get_access_token()
     except ClientError as err:
         raise ConfigEntryNotReady from err
-    coordinator = AutomowerDataUpdateCoordinator(hass, automower_api)
+    coordinator = AutomowerDataUpdateCoordinator(hass, automower_api, entry)
     await coordinator.async_config_entry_first_refresh()
+    entry.async_create_background_task(
+        hass,
+        coordinator.client_listen(hass, entry, automower_api),
+        "websocket_task",
+    )
 
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
-    entry.async_on_unload(
-        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, coordinator.shutdown)
-    )
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     return True
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Handle unload of an entry."""
-    coordinator: AutomowerDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    await coordinator.shutdown()
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
     if unload_ok:
         hass.data[DOMAIN].pop(entry.entry_id)

homeassistant/components/husqvarna_automower/coordinator.py

@@ -1,23 +1,28 @@
 """Data UpdateCoordinator for the Husqvarna Automower integration."""
+import asyncio
 from datetime import timedelta
 import logging
-from typing import Any
 
+from aioautomower.exceptions import ApiException, HusqvarnaWSServerHandshakeError
 from aioautomower.model import MowerAttributes
+from aioautomower.session import AutomowerSession
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .api import AsyncConfigEntryAuth
 from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
+MAX_WS_RECONNECT_TIME = 600
 
 
 class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[dict[str, MowerAttributes]]):
     """Class to manage fetching Husqvarna data."""
 
-    def __init__(self, hass: HomeAssistant, api: AsyncConfigEntryAuth) -> None:
+    def __init__(
+        self, hass: HomeAssistant, api: AutomowerSession, entry: ConfigEntry
+    ) -> None:
         """Initialize data updater."""
         super().__init__(
             hass,
@@ -35,13 +40,39 @@ async def _async_update_data(self) -> dict[str, MowerAttributes]:
             await self.api.connect()
             self.api.register_data_callback(self.callback)
             self.ws_connected = True
-        return await self.api.get_status()
-
-    async def shutdown(self, *_: Any) -> None:
-        """Close resources."""
-        await self.api.close()
+        try:
+            return await self.api.get_status()
+        except ApiException as err:
+            raise UpdateFailed(err) from err
 
     @callback
     def callback(self, ws_data: dict[str, MowerAttributes]) -> None:
         """Process websocket callbacks and write them to the DataUpdateCoordinator."""
         self.async_set_updated_data(ws_data)
+
+    async def client_listen(
+        self,
+        hass: HomeAssistant,
+        entry: ConfigEntry,
+        automower_client: AutomowerSession,
+        reconnect_time: int = 2,
+    ) -> None:
+        """Listen with the client."""
+        try:
+            await automower_client.auth.websocket_connect()
+            reconnect_time = 2
+            await automower_client.start_listening()
+        except HusqvarnaWSServerHandshakeError as err:
+            _LOGGER.debug(
+                "Failed to connect to websocket. Trying to reconnect: %s", err
+            )
+
+        if not hass.is_stopping:
+            await asyncio.sleep(reconnect_time)
+            reconnect_time = min(reconnect_time * 2, MAX_WS_RECONNECT_TIME)
+            await self.client_listen(
+                hass=hass,
+                entry=entry,
+                automower_client=automower_client,
+                reconnect_time=reconnect_time,
+            )

homeassistant/components/husqvarna_automower/manifest.json

@@ -6,5 +6,5 @@
   "dependencies": ["application_credentials"],
   "documentation": "https://www.home-assistant.io/integrations/husqvarna_automower",
   "iot_class": "cloud_push",
-  "requirements": ["aioautomower==2024.2.7"]
+  "requirements": ["aioautomower==2024.2.10"]
 }

requirements_all.txt

@@ -206,7 +206,7 @@ aioaseko==0.0.2
 aioasuswrt==1.4.0
 
 # homeassistant.components.husqvarna_automower
-aioautomower==2024.2.7
+aioautomower==2024.2.10
 
 # homeassistant.components.azure_devops
 aioazuredevops==1.3.5

requirements_test_all.txt

@@ -185,7 +185,7 @@ aioaseko==0.0.2
 aioasuswrt==1.4.0
 
 # homeassistant.components.husqvarna_automower
-aioautomower==2024.2.7
+aioautomower==2024.2.10
 
 # homeassistant.components.azure_devops
 aioazuredevops==1.3.5

tests/components/husqvarna_automower/conftest.py

@@ -4,6 +4,7 @@
 from unittest.mock import AsyncMock, patch
 
 from aioautomower.utils import mower_list_to_dictionary_dataclass
+from aiohttp import ClientWebSocketResponse
 import pytest
 
 from homeassistant.components.application_credentials import (
@@ -82,4 +83,11 @@ def mock_automower_client() -> Generator[AsyncMock, None, None]:
         client.get_status.return_value = mower_list_to_dictionary_dataclass(
             load_json_value_fixture("mower.json", DOMAIN)
         )
+
+        async def websocket_connect() -> ClientWebSocketResponse:
+            """Mock listen."""
+            return ClientWebSocketResponse
+
+        client.auth = AsyncMock(side_effect=websocket_connect)
+
         yield client

tests/components/husqvarna_automower/test_init.py

@@ -1,8 +1,11 @@
 """Tests for init module."""
+from datetime import timedelta
 import http
 import time
 from unittest.mock import AsyncMock
 
+from aioautomower.exceptions import ApiException, HusqvarnaWSServerHandshakeError
+from freezegun.api import FrozenDateTimeFactory
 import pytest
 
 from homeassistant.components.husqvarna_automower.const import DOMAIN, OAUTH2_TOKEN
@@ -11,7 +14,7 @@
 
 from . import setup_integration
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_fire_time_changed
 from tests.test_util.aiohttp import AiohttpClientMocker
 
 
@@ -66,3 +69,42 @@ async def test_expired_token_refresh_failure(
     await setup_integration(hass, mock_config_entry)
 
     assert mock_config_entry.state is expected_state
+
+
+async def test_update_failed(
+    hass: HomeAssistant,
+    mock_automower_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test load and unload entry."""
+    getattr(mock_automower_client, "get_status").side_effect = ApiException(
+        "Test error"
+    )
+    await setup_integration(hass, mock_config_entry)
+    entry = hass.config_entries.async_entries(DOMAIN)[0]
+
+    assert entry.state == ConfigEntryState.SETUP_RETRY
+
+
+async def test_websocket_not_available(
+    hass: HomeAssistant,
+    mock_automower_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    caplog: pytest.LogCaptureFixture,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test trying reload the websocket."""
+    mock_automower_client.start_listening.side_effect = HusqvarnaWSServerHandshakeError(
+        "Boom"
+    )
+    await setup_integration(hass, mock_config_entry)
+    assert "Failed to connect to websocket. Trying to reconnect: Boom" in caplog.text
+    assert mock_automower_client.auth.websocket_connect.call_count == 1
+    assert mock_automower_client.start_listening.call_count == 1
+    assert mock_config_entry.state == ConfigEntryState.LOADED
+    freezer.tick(timedelta(seconds=2))
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    assert mock_automower_client.auth.websocket_connect.call_count == 2
+    assert mock_automower_client.start_listening.call_count == 2
+    assert mock_config_entry.state == ConfigEntryState.LOADED