websocket.py

Go to the documentation of this file.

"""Websocket handler for UniFi Network integration."""
 
from __future__ import annotations
 
import asyncio
from datetime import datetime, timedelta
 
import aiohttp
import aiounifi
 
from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
from homeassistant.helpers.dispatcher import async_dispatcher_send
from homeassistant.helpers.event import async_track_time_interval
 
from ..const import LOGGER
 
RETRY_TIMER = 15
CHECK_WEBSOCKET_INTERVAL = timedelta(minutes=1)
 
 
class UnifiWebsocket:
    """Manages a single UniFi Network instance."""
 
    def __init__(
        self, hass: HomeAssistant, api: aiounifi.Controller, signal: str
    ) -> None:
        """Initialize the system."""
        self.hasshass = hass
        self.apiapi = api
        self.signalsignal = signal
 
        self.ws_taskws_task: asyncio.Task | None = None
        self._cancel_websocket_check_cancel_websocket_check: CALLBACK_TYPE | None = None
 
        self.availableavailable = True
 
    @callback
    def start(self) -> None:
        """Start websocket handler."""
        self._cancel_websocket_check_cancel_websocket_check = async_track_time_interval(
            self.hasshass, self._async_watch_websocket_async_watch_websocket, CHECK_WEBSOCKET_INTERVAL
        )
        self.start_websocketstart_websocket()
 
    @callback
    def stop(self) -> None:
        """Stop websocket handler."""
        if self._cancel_websocket_check_cancel_websocket_check:
            self._cancel_websocket_check_cancel_websocket_check()
            self._cancel_websocket_check_cancel_websocket_check = None
 
        if self.ws_taskws_task is not None:
            self.ws_taskws_task.cancel()
 
    async def stop_and_wait(self) -> None:
        """Stop websocket handler and await tasks."""
        if self._cancel_websocket_check_cancel_websocket_check:
            self._cancel_websocket_check_cancel_websocket_check()
            self._cancel_websocket_check_cancel_websocket_check = None
 
        if self.ws_taskws_task is not None:
            self.stopstop()
 
            _, pending = await asyncio.wait([self.ws_taskws_task], timeout=10)
 
            if pending:
                LOGGER.warning(
                    "Unloading UniFi Network (%s). Task %s did not complete in time",
                    self.apiapi.connectivity.config.host,
                    self.ws_taskws_task,
                )
 
    @callback
    def start_websocket(self) -> None:
        """Start up connection to websocket."""
 
        async def _websocket_runner() -> None:
            """Start websocket."""
            try:
                await self.apiapi.start_websocket()
            except (aiohttp.ClientConnectorError, aiohttp.WSServerHandshakeError):
                LOGGER.error("Websocket setup failed")
            except aiounifi.WebsocketError:
                LOGGER.error("Websocket disconnected")
 
            self.availableavailable = False
            async_dispatcher_send(self.hasshass, self.signalsignal)
            self.hasshass.loop.call_later(RETRY_TIMER, self.reconnectreconnect, True)
 
        if not self.availableavailable:
            self.availableavailable = True
            async_dispatcher_send(self.hasshass, self.signalsignal)
 
        self.ws_taskws_task = self.hasshass.loop.create_task(_websocket_runner())
 
    @callback
    def reconnect(self, log: bool = False) -> None:
        """Prepare to reconnect UniFi session."""
 
        async def _reconnect() -> None:
            """Try to reconnect UniFi Network session."""
            try:
                async with asyncio.timeout(5):
                    await self.apiapi.login()
 
            except (
                TimeoutError,
                aiounifi.BadGateway,
                aiounifi.ServiceUnavailable,
                aiounifi.AiounifiException,
            ) as exc:
                LOGGER.debug("Schedule reconnect to UniFi Network '%s'", exc)
                self.hasshass.loop.call_later(RETRY_TIMER, self.reconnectreconnect)
 
            else:
                self.start_websocketstart_websocket()
 
        if log:
            LOGGER.info("Will try to reconnect to UniFi Network")
 
        self.hasshass.loop.create_task(_reconnect())
 
    @callback
    def _async_watch_websocket(self, now: datetime) -> None:
        """Watch timestamp for last received websocket message."""
        LOGGER.debug(
            "Last received websocket timestamp: %s",
            self.apiapi.connectivity.ws_message_received,
        )