diff --git a/homeassistant/components/husqvarna_automower/__init__.py b/homeassistant/components/husqvarna_automower/__init__.py index 7ed8a6b23e8dae..20218229385f77 100644 --- a/homeassistant/components/husqvarna_automower/__init__.py +++ b/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,13 +37,15 @@ 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 @@ -52,8 +53,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: 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) diff --git a/homeassistant/components/husqvarna_automower/coordinator.py b/homeassistant/components/husqvarna_automower/coordinator.py index 70d69f90549b9c..2840823415aae3 100644 --- a/homeassistant/components/husqvarna_automower/coordinator.py +++ b/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, + ) diff --git a/homeassistant/components/husqvarna_automower/manifest.json b/homeassistant/components/husqvarna_automower/manifest.json index 1eb40bfad33459..dc40116f31efcb 100644 --- a/homeassistant/components/husqvarna_automower/manifest.json +++ b/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"] } diff --git a/requirements_all.txt b/requirements_all.txt index dd6005ccd83c13..95d6c0f5be70d4 100644 --- a/requirements_all.txt +++ b/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 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 20925be15c5578..c5f3e921b6fd1c 100644 --- a/requirements_test_all.txt +++ b/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 diff --git a/tests/components/husqvarna_automower/conftest.py b/tests/components/husqvarna_automower/conftest.py index 89c0133cd0b495..3194f1b3188b10 100644 --- a/tests/components/husqvarna_automower/conftest.py +++ b/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 diff --git a/tests/components/husqvarna_automower/test_init.py b/tests/components/husqvarna_automower/test_init.py index 14460ad5d21f5e..c11e4ac4cc7d3b 100644 --- a/tests/components/husqvarna_automower/test_init.py +++ b/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