ha_gehome/ge_kitchen/update_coordinator.py

"""Data update coordinator for GE Kitchen Appliances"""

import asyncio
import logging
from typing import Any, Dict, Iterable, Optional, Tuple

from gekitchensdk import (
    EVENT_APPLIANCE_INITIAL_UPDATE,
    EVENT_APPLIANCE_UPDATE_RECEIVED,
    EVENT_CONNECTED,
    EVENT_DISCONNECTED,
    EVENT_GOT_APPLIANCE_LIST,
    ErdCodeType,
    GeAppliance,
    GeWebsocketClient,
)

from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
from homeassistant.core import HomeAssistant
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator

from .const import DOMAIN, EVENT_ALL_APPLIANCES_READY, UPDATE_INTERVAL
from .devices import ApplianceApi, get_appliance_api_type

_LOGGER = logging.getLogger(__name__)

class GeKitchenUpdateCoordinator(DataUpdateCoordinator):
    """Define a wrapper class to update GE Kitchen data."""

    def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry) -> None:
        """Set up the GeKitchenUpdateCoordinator class."""
        self._hass = hass
        self._config_entry = config_entry
        self._username = config_entry.data[CONF_USERNAME]
        self._password = config_entry.data[CONF_PASSWORD]
        self.client = None  # type: Optional[GeWebsocketClient]
        self._appliance_apis = {}  # type: Dict[str, ApplianceApi]

        # Some record keeping to let us know when we can start generating entities
        self._got_roster = False
        self._init_done = False
        self.initialization_future = asyncio.Future()

        super().__init__(hass, _LOGGER, name=DOMAIN)

    def create_ge_client(self, event_loop: Optional[asyncio.AbstractEventLoop]) -> GeWebsocketClient:
        """
        Create a new GeClient object with some helpful callbacks.

        :param event_loop: Event loop
        :return: GeWebsocketClient
        """
        client = GeWebsocketClient(event_loop=event_loop, username=self._username, password=self._password)
        client.add_event_handler(EVENT_APPLIANCE_INITIAL_UPDATE, self.on_device_initial_update)
        client.add_event_handler(EVENT_APPLIANCE_UPDATE_RECEIVED, self.on_device_update)
        client.add_event_handler(EVENT_GOT_APPLIANCE_LIST, self.on_appliance_list)
        client.add_event_handler(EVENT_DISCONNECTED, self.on_disconnect)
        client.add_event_handler(EVENT_CONNECTED, self.on_connect)
        return client

    @property
    def appliances(self) -> Iterable[GeAppliance]:
        return self.client.appliances.values()

    @property
    def appliance_apis(self) -> Dict[str, ApplianceApi]:
        return self._appliance_apis

    def _get_appliance_api(self, appliance: GeAppliance) -> ApplianceApi:
        api_type = get_appliance_api_type(appliance.appliance_type)
        return api_type(self, appliance)

    def regenerate_appliance_apis(self):
        """Regenerate the appliance_apis dictionary, adding elements as necessary."""
        for jid, appliance in self.client.appliances.keys():
            if jid not in self._appliance_apis:
                self._appliance_apis[jid] = self._get_appliance_api(appliance)

    def maybe_add_appliance_api(self, appliance: GeAppliance):
        mac_addr = appliance.mac_addr
        if mac_addr not in self.appliance_apis:
            _LOGGER.debug(f"Adding appliance api for appliance {mac_addr} ({appliance.appliance_type})")
            api = self._get_appliance_api(appliance)
            api.build_entities_list()
            self.appliance_apis[mac_addr] = api

    async def get_client(self) -> GeWebsocketClient:
        """Get a new GE Websocket client."""
        if self.client is not None:
            await self.client.disconnect()

        loop = self._hass.loop
        self.client = self.create_ge_client(event_loop=loop)
        return self.client

    async def async_start_client(self):
        """Start a new GeClient in the HASS event loop."""
        _LOGGER.debug('Running client')
        client = await self.get_client()

        session = self._hass.helpers.aiohttp_client.async_get_clientsession()
        await client.async_get_credentials(session)
        fut = asyncio.ensure_future(client.async_run_client(), loop=self._hass.loop)
        _LOGGER.debug('Client running')
        return fut

    async def _kill_client(self):
        """Kill the client.  Leaving this in for testing purposes."""
        await asyncio.sleep(30)
        _LOGGER.critical('Killing the connection.  Popcorn time.')
        await self.client.websocket.close()

    async def on_device_update(self, data: Tuple[GeAppliance, Dict[ErdCodeType, Any]]):
        """Let HA know there's new state."""
        self.last_update_success = True
        appliance, _ = data
        try:
            api = self.appliance_apis[appliance.mac_addr]
        except KeyError:
            return
        for entity in api.entities:
            _LOGGER.debug(f'Updating {entity} ({entity.unique_id}, {entity.entity_id})')
            entity.async_write_ha_state()

    @property
    def all_appliances_updated(self) -> bool:
        """True if all appliances have had an initial update."""
        return all([a.initialized for a in self.appliances])

    async def on_appliance_list(self, _):
        """When we get an appliance list, mark it and maybe trigger all ready."""
        _LOGGER.debug('Got roster update')
        self.last_update_success = True
        if not self._got_roster:
            self._got_roster = True
            await asyncio.sleep(5)  # After the initial roster update, wait a bit and hit go
            await self.async_maybe_trigger_all_ready()

    async def on_device_initial_update(self, appliance: GeAppliance):
        """When an appliance first becomes ready, let the system know and schedule periodic updates."""
        _LOGGER.debug(f'Got initial update for {appliance.mac_addr}')
        self.last_update_success = True
        self.maybe_add_appliance_api(appliance)
        await self.async_maybe_trigger_all_ready()
        _LOGGER.debug(f'Requesting updates for {appliance.mac_addr}')
        while not self.client.websocket.closed and appliance.available:
            await asyncio.sleep(UPDATE_INTERVAL)
            await appliance.async_request_update()
        _LOGGER.debug(f'No longer requesting updates for {appliance.mac_addr}')

    async def on_disconnect(self, _):
        """Handle disconnection."""
        _LOGGER.debug("Disconnected. Attempting to reconnect.")
        self.last_update_success = False

        flow_context = {
            "source": "reauth",
            "unique_id": self._config_entry.unique_id,
        }

        matching_flows = [
            flow
            for flow in self.hass.config_entries.flow.async_progress()
            if flow["context"] == flow_context
        ]

        if not matching_flows:
            self.hass.async_create_task(
                self.hass.config_entries.flow.async_init(
                    DOMAIN, context=flow_context, data=self._config_entry.data,
                )
            )

    async def on_connect(self, _):
        """Set state upon connection."""
        self.last_update_success = True

    async def async_maybe_trigger_all_ready(self):
        """See if we're all ready to go, and if so, let the games begin."""
        if self._init_done or self.initialization_future.done():
            # Been here, done this
            return
        if self._got_roster and self.all_appliances_updated:
            _LOGGER.debug('Ready to go.  Waiting 2 seconds and setting init future result.')
            # The the flag and wait to prevent two different fun race conditions
            self._init_done = True
            await asyncio.sleep(2)
            self.initialization_future.set_result(True)
            await self.client.async_event(EVENT_ALL_APPLIANCES_READY, None)