mirror of https://github.com/home-assistant/core
48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
"""Data UpdateCoordinator for the Husqvarna Automower integration."""
|
|
from datetime import timedelta
|
|
import logging
|
|
from typing import Any
|
|
|
|
from aioautomower.model import MowerAttributes, MowerList
|
|
|
|
from homeassistant.core import HomeAssistant, callback
|
|
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
|
|
|
|
from .api import AsyncConfigEntryAuth
|
|
from .const import DOMAIN
|
|
|
|
_LOGGER = logging.getLogger(__name__)
|
|
|
|
|
|
class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[dict[str, MowerAttributes]]):
|
|
"""Class to manage fetching Husqvarna data."""
|
|
|
|
def __init__(self, hass: HomeAssistant, api: AsyncConfigEntryAuth) -> None:
|
|
"""Initialize data updater."""
|
|
super().__init__(
|
|
hass,
|
|
_LOGGER,
|
|
name=DOMAIN,
|
|
update_interval=timedelta(minutes=5),
|
|
)
|
|
self.api = api
|
|
|
|
self.ws_connected: bool = False
|
|
|
|
async def _async_update_data(self) -> dict[str, MowerAttributes]:
|
|
"""Subscribe for websocket and poll data from the API."""
|
|
if not self.ws_connected:
|
|
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()
|
|
|
|
@callback
|
|
def callback(self, ws_data: MowerList) -> None:
|
|
"""Process websocket callbacks and write them to the DataUpdateCoordinator."""
|
|
self.async_set_updated_data(ws_data)
|