restore_state.RestoreStateData.async_get_instance is deprecated about ha-tplink-deco HOT 3 CLOSED

"""
Custom integration to integrate TP-Link Deco with Home Assistant.

For more details about this integration, please refer to
https://github.com/amosyuen/ha-tplink-deco
"""
import asyncio
import logging
from datetime import timedelta
from typing import Any
from typing import cast

import homeassistant.helpers.config_validation as cv
import voluptuous as vol
from homeassistant.components.device_tracker.const import CONF_CONSIDER_HOME
from homeassistant.components.device_tracker.const import CONF_SCAN_INTERVAL
from homeassistant.components.device_tracker.const import (
    DOMAIN as DEVICE_TRACKER_DOMAIN,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_DEVICE_ID
from homeassistant.const import CONF_HOST
from homeassistant.const import CONF_PASSWORD
from homeassistant.const import CONF_USERNAME
from homeassistant.core import Config
from homeassistant.core import HomeAssistant
from homeassistant.core import ServiceCall
from homeassistant.helpers import device_registry
from homeassistant.helpers import entity_registry
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.helpers.restore_state import RestoreEntity

from .api import TplinkDecoApi
from .const import ATTR_DEVICE_TYPE
from .const import CONF_TIMEOUT_ERROR_RETRIES
from .const import CONF_TIMEOUT_SECONDS
from .const import CONF_VERIFY_SSL
from .const import COORDINATOR_CLIENTS_KEY
from .const import COORDINATOR_DECOS_KEY
from .const import DEFAULT_CONSIDER_HOME
from .const import DEFAULT_SCAN_INTERVAL
from .const import DEFAULT_TIMEOUT_ERROR_RETRIES
from .const import DEFAULT_TIMEOUT_SECONDS
from .const import DEVICE_TYPE_DECO
from .const import DOMAIN
from .const import PLATFORMS
from .const import SERVICE_REBOOT_DECO
from .coordinator import TpLinkDeco
from .coordinator import TpLinkDecoClient
from .coordinator import TplinkDecoClientUpdateCoordinator
from .coordinator import TpLinkDecoData
from .coordinator import TplinkDecoUpdateCoordinator

_LOGGER: logging.Logger = logging.getLogger(__name__)


async def async_create_and_refresh_coordinators(
    hass: HomeAssistant,
    config_data: dict[str:Any],
    consider_home_seconds,
    update_interval: timedelta = None,
    deco_data: TpLinkDecoData = None,
    client_data: dict[str:TpLinkDecoClient] = None,
):
    host = config_data.get(CONF_HOST)
    username = config_data.get(CONF_USERNAME)
    password = config_data.get(CONF_PASSWORD)
    timeout_error_retries = config_data.get(CONF_TIMEOUT_ERROR_RETRIES)
    timeout_seconds = config_data.get(CONF_TIMEOUT_SECONDS)
    verify_ssl = config_data.get(CONF_VERIFY_SSL)
    session = async_get_clientsession(hass)

    api = TplinkDecoApi(
        session,
        host,
        username,
        password,
        verify_ssl,
        timeout_error_retries,
        timeout_seconds,
    )
    deco_coordinator = TplinkDecoUpdateCoordinator(
        hass, api, update_interval, deco_data
    )
    await deco_coordinator.async_config_entry_first_refresh()
    clients_coordinator = TplinkDecoClientUpdateCoordinator(
        hass,
        api,
        deco_coordinator,
        consider_home_seconds,
        update_interval,
        client_data,
    )
    await clients_coordinator.async_config_entry_first_refresh()

    return {
        COORDINATOR_DECOS_KEY: deco_coordinator,
        COORDINATOR_CLIENTS_KEY: clients_coordinator,
    }


async def async_create_config_data(hass: HomeAssistant, config_entry: ConfigEntry):
    consider_home_seconds = config_entry.data.get(
        CONF_CONSIDER_HOME, DEFAULT_CONSIDER_HOME
    )
    scan_interval_seconds = config_entry.data.get(
        CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL
    )
    update_interval = timedelta(seconds=scan_interval_seconds)

    # Load tracked entities from registry
    existing_entries = entity_registry.async_entries_for_config_entry(
        entity_registry.async_get(hass),
        config_entry.entry_id,
    )
    deco_data = TpLinkDecoData()
    client_data = {}

    last_states = hass.states.async_all()
    for entry in existing_entries:
        if entry.domain != DEVICE_TRACKER_DOMAIN:
            continue
        state = hass.states.get(entry.entity_id)
        if state is None:
            continue
        device_type = state.attributes.get(ATTR_DEVICE_TYPE)
        if device_type is None:
            continue
        if device_type == DEVICE_TYPE_DECO:
            deco = TpLinkDeco(entry.unique_id)
            deco.name = entry.original_name
            deco_data.decos[entry.unique_id] = deco
        else:
            client = TpLinkDecoClient(entry.unique_id)
            client.name = entry.original_name
            client_data[entry.unique_id] = client

    return await async_create_and_refresh_coordinators(
        hass,
        config_entry.data,
        consider_home_seconds,
        update_interval,
        deco_data,
        client_data,
    )


async def async_setup(hass: HomeAssistant, config: Config):
    """Set up this integration using YAML is not supported."""
    return True


async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry):
    """Set up this integration using UI."""
    if hass.data.get(DOMAIN) is None:
        hass.data.setdefault(DOMAIN, {})

    data = await async_create_config_data(hass, config_entry)
    hass.data[DOMAIN][config_entry.entry_id] = data
    deco_coordinator = data[COORDINATOR_DECOS_KEY]

    for platform in PLATFORMS:
        hass.async_add_job(
            hass.config_entries.async_forward_entry_setup(config_entry, platform)
        )

    async def async_reboot_deco(service: ServiceCall) -> None:
        dr = device_registry.async_get(hass=hass)
        device_ids = cast([str], service.data.get(ATTR_DEVICE_ID))
        macs = []
        for device_id in device_ids:
            device = dr.async_get(device_id)
            if device is None:
                raise Exception(f"Device ID {device_id} is not a TP-Link Deco device")
            ids = device.identifiers
            id = next(iter(ids)) if len(ids) == 1 else None
            if id[0] != DOMAIN:
                raise Exception(
                    f"Device ID {device_id} does not have {DOMAIN} MAC identifier"
                )
            macs.append(id[1])
        await deco_coordinator.api.async_reboot_decos(macs)

    hass.services.async_register(
        DOMAIN,
        SERVICE_REBOOT_DECO,
        async_reboot_deco,
        schema=vol.Schema(
            {
                vol.Required(ATTR_DEVICE_ID): vol.All(cv.ensure_list(str), [str]),
            }
        ),
    )

    config_entry.async_on_unload(config_entry.add_update_listener(update_listener))
    return True


async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
    """Handle removal of an entry."""
    data = hass.data[DOMAIN][config_entry.entry_id]
    deco_coordinator = data.get(COORDINATOR_DECOS_KEY)
    clients_coordinator = data.get(COORDINATOR_CLIENTS_KEY)
    if deco_coordinator is not None:
        await deco_coordinator.async_stop()
    if clients_coordinator is not None:
        await clients_coordinator.async_stop()
    unload_ok = await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS)
    if unload_ok:
        hass.data[DOMAIN].pop(config_entry.entry_id)
    return unload_ok


async def update_listener(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
    """Handle options update."""
    data = hass.data[DOMAIN][config_entry.entry_id]
    deco_coordinator = data[COORDINATOR_DECOS_KEY]
    clients_coordinator = data[COORDINATOR_CLIENTS_KEY]

    await deco_coordinator.api.async_update_options(config_entry.data)
    deco_coordinator.update_interval = timedelta(
        seconds=config_entry.data.get(CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL)
    )

    await deco_coordinator.async_config_entry_first_refresh()
    await clients_coordinator.async_config_entry_first_refresh()


class TpLinkDecoDevice(RestoreEntity):
    """Representation of a TP-Link Deco device."""

    def __init__(
        self,
        config_entry: ConfigEntry,
        coordinator: TplinkDecoUpdateCoordinator,
        deco_id: str,
    ):
        """Initialize the TP-Link Deco device."""
        self.config_entry = config_entry
        self.coordinator = coordinator
        self.deco_id = deco_id
        self.deco: TpLinkDeco | None = None

    @property
    def unique_id(self) -> str:
        """Return a unique ID for the device."""
        return f"{DOMAIN}-{self.deco_id}"

    @property
    def name(self) -> str:
        """Return the name of the device."""
        if self.deco is not None:
            return self.deco.name
        return ""

    @property
    def available(self) -> bool:
        """Return True if the device is available."""
        return self.coordinator.last_update_success and self.deco_id in self.coordinator.data.decos

    async def async_added_to_hass(self):
        """Subscribe to updates."""
        self.async_on_remove(
            self.coordinator.async_add_listener(self.async_write_ha_state)
        )
        await super().async_added_to_hass()

    async def async_update(self):
        """Update the device."""
        self.deco = self.coordinator.data.decos.get(self.deco_id)

    async def async_turn_on(self, **kwargs: Any) -> None:
        """Turn on the device."""
        await self.coordinator.api.async_turn_on_deco(self.deco_id)

    async def async_turn_off(self, **kwargs: Any) -> None:
        """Turn off the device."""
        await self.coordinator.api.async_turn_off_deco(self.deco_id)


class TpLinkDecoClientDevice(RestoreEntity):
    """Representation of a TP-Link Deco client device."""

    def __init__(
        self,
        config_entry: ConfigEntry,
        coordinator: TplinkDecoClientUpdateCoordinator,
        client_id: str,
    ):
        """Initialize the TP-Link Deco client device."""
        self.config_entry = config_entry
        self.coordinator = coordinator
        self.client_id = client_id
        self.client: TpLinkDecoClient | None = None

    @property
    def unique_id(self) -> str:
        """Return a unique ID for the device."""
        return f"{DOMAIN}-{self.client_id}"

    @property
    def name(self) -> str:
        """Return the name of the device."""
        if self.client is not None:
            return self.client.name
        return ""

    @property
    def available(self) -> bool:
        """Return True if the device is available."""
        return (
            self.coordinator.last_update_success
            and self.client_id in self.coordinator.data.clients
        )

    @property
    def device_info(self) -> dict[str, Any]:
        """Return the device info."""
        device_id = self.client_id.split(":")[0]
        device = device_registry.async_get(hass=self.hass).async_get(device_id)
        if device is not None:
            return {
                "identifiers": {(DOMAIN, device_id)},
                "name": device.name,
                "manufacturer": "TP-Link",
                "model": device.model,
            }
        return {}

    @property
    def device_state_attributes(self) -> dict[str, Any]:
        """Return the state attributes of the device."""
        if self.client is not None:
            return {
                ATTR_DEVICE_TYPE: self.client.device_type,
            }
        return {}

    async def async_added_to_hass(self):
        """Subscribe to updates."""
        self.async_on_remove(
            self.coordinator.async_add_listener(self.async_write_ha_state)
        )
        await super().async_added_to_hass()

    async def async_update(self):
        """Update the device."""
        self.client = self.coordinator.data.clients.get(self.client_id)

    async def async_turn_on(self, **kwargs: Any) -> None:
        """Turn on the device."""
        await self.coordinator.api.async_turn_on_client(self.client_id)

    async def async_turn_off(self, **kwargs: Any) -> None:
        """Turn off the device."""
        await self.coordinator.api.async_turn_off_client(self.client_id)


async def async_migrate_entry(hass, config_entry):
    """Migrate old entry."""
    _LOGGER.debug("Migrating TP-Link Deco entry from version %s", config_entry.version)
    data = config_entry.data
    options = config_entry.options

    version = config_entry.version

    if version == 1:
        data[CONF_VERIFY_SSL] = True
        options[CONF_SCAN_INTERVAL] = options.pop(CONF_UPDATE_INTERVAL, 5 * 60)

        config_entry.version = 2
        hass.config_entries.async_update_entry(config_entry, data=data, options=options)

    _LOGGER.info("Migration to version %s successful", config_entry.version)
    return True