core/homeassistant/components/zwave_js/update.py

"""Representation of Z-Wave updates."""
from __future__ import annotations

import asyncio
from collections.abc import Callable
from datetime import datetime, timedelta
from math import floor
from typing import Any

from awesomeversion import AwesomeVersion
from zwave_js_server.client import Client as ZwaveClient
from zwave_js_server.const import NodeStatus
from zwave_js_server.exceptions import BaseZwaveJSServerError, FailedZWaveCommand
from zwave_js_server.model.driver import Driver
from zwave_js_server.model.firmware import (
    FirmwareUpdateFinished,
    FirmwareUpdateInfo,
    FirmwareUpdateProgress,
    FirmwareUpdateStatus,
)
from zwave_js_server.model.node import Node as ZwaveNode

from homeassistant.components.update import (
    UpdateDeviceClass,
    UpdateEntity,
    UpdateEntityFeature,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import EntityCategory
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.event import async_call_later
from homeassistant.helpers.start import async_at_start

from .const import API_KEY_FIRMWARE_UPDATE_SERVICE, DATA_CLIENT, DOMAIN, LOGGER
from .helpers import get_device_info, get_valueless_base_unique_id

PARALLEL_UPDATES = 1


async def async_setup_entry(
    hass: HomeAssistant,
    config_entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up Z-Wave button from config entry."""
    client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT]

    semaphore = asyncio.Semaphore(3)

    @callback
    def async_add_firmware_update_entity(node: ZwaveNode) -> None:
        """Add firmware update entity."""
        driver = client.driver
        assert driver is not None  # Driver is ready before platforms are loaded.
        async_add_entities([ZWaveNodeFirmwareUpdate(driver, node, semaphore)])

    config_entry.async_on_unload(
        async_dispatcher_connect(
            hass,
            f"{DOMAIN}_{config_entry.entry_id}_add_firmware_update_entity",
            async_add_firmware_update_entity,
        )
    )


class ZWaveNodeFirmwareUpdate(UpdateEntity):
    """Representation of a firmware update entity."""

    _attr_entity_category = EntityCategory.CONFIG
    _attr_device_class = UpdateDeviceClass.FIRMWARE
    _attr_supported_features = (
        UpdateEntityFeature.INSTALL
        | UpdateEntityFeature.RELEASE_NOTES
        | UpdateEntityFeature.PROGRESS
    )
    _attr_has_entity_name = True
    _attr_should_poll = False

    def __init__(
        self, driver: Driver, node: ZwaveNode, semaphore: asyncio.Semaphore
    ) -> None:
        """Initialize a Z-Wave device firmware update entity."""
        self.driver = driver
        self.node = node
        self.semaphore = semaphore
        self._latest_version_firmware: FirmwareUpdateInfo | None = None
        self._status_unsub: Callable[[], None] | None = None
        self._poll_unsub: Callable[[], None] | None = None
        self._progress_unsub: Callable[[], None] | None = None
        self._finished_unsub: Callable[[], None] | None = None
        self._num_files_installed: int = 0
        self._finished_event = asyncio.Event()
        self._finished_status: FirmwareUpdateStatus | None = None

        # Entity class attributes
        self._attr_name = "Firmware"
        self._base_unique_id = get_valueless_base_unique_id(driver, node)
        self._attr_unique_id = f"{self._base_unique_id}.firmware_update"
        self._attr_installed_version = self._attr_latest_version = node.firmware_version
        # device may not be precreated in main handler yet
        self._attr_device_info = get_device_info(driver, node)

    @callback
    def _update_on_status_change(self, _: dict[str, Any]) -> None:
        """Update the entity when node is awake."""
        self._status_unsub = None
        self.hass.async_create_task(self._async_update())

    @callback
    def _update_progress(self, event: dict[str, Any]) -> None:
        """Update install progress on event."""
        progress: FirmwareUpdateProgress = event["firmware_update_progress"]
        if not self._latest_version_firmware:
            return
        # We will assume that each file in the firmware update represents an equal
        # percentage of the overall progress. This is likely not true because each file
        # may be a different size, but it's the best we can do since we don't know the
        # total number of fragments across all files.
        self._attr_in_progress = floor(
            100
            * (
                self._num_files_installed
                + (progress.sent_fragments / progress.total_fragments)
            )
            / len(self._latest_version_firmware.files)
        )
        self.async_write_ha_state()

    @callback
    def _update_finished(self, event: dict[str, Any]) -> None:
        """Update install progress on event."""
        finished: FirmwareUpdateFinished = event["firmware_update_finished"]
        self._finished_status = finished.status
        self._finished_event.set()

    @callback
    def _unsub_firmware_events_and_reset_progress(
        self, write_state: bool = False
    ) -> None:
        """Unsubscribe from firmware events and reset update install progress."""
        if self._progress_unsub:
            self._progress_unsub()
            self._progress_unsub = None

        if self._finished_unsub:
            self._finished_unsub()
            self._finished_unsub = None

        self._finished_status = None
        self._finished_event.clear()
        self._num_files_installed = 0
        self._attr_in_progress = 0
        if write_state:
            self.async_write_ha_state()

    async def _async_update(self, _: HomeAssistant | datetime | None = None) -> None:
        """Update the entity."""
        self._poll_unsub = None

        # If device is asleep/dead, wait for it to wake up/become alive before
        # attempting an update
        for status, event_name in (
            (NodeStatus.ASLEEP, "wake up"),
            (NodeStatus.DEAD, "alive"),
        ):
            if self.node.status == status:
                if not self._status_unsub:
                    self._status_unsub = self.node.once(
                        event_name, self._update_on_status_change
                    )
                return

        try:
            async with self.semaphore:
                available_firmware_updates = (
                    await self.driver.controller.async_get_available_firmware_updates(
                        self.node, API_KEY_FIRMWARE_UPDATE_SERVICE
                    )
                )
        except FailedZWaveCommand as err:
            LOGGER.debug(
                "Failed to get firmware updates for node %s: %s",
                self.node.node_id,
                err,
            )
        else:
            if available_firmware_updates:
                self._latest_version_firmware = latest_firmware = max(
                    available_firmware_updates,
                    key=lambda x: AwesomeVersion(x.version),
                )

                # If we have an available firmware update that is a higher version than
                # what's on the node, we should advertise it, otherwise there is
                # nothing to do.
                new_version = latest_firmware.version
                current_version = self.node.firmware_version
                if AwesomeVersion(new_version) > AwesomeVersion(current_version):
                    self._attr_latest_version = new_version
                    self.async_write_ha_state()
        finally:
            self._poll_unsub = async_call_later(
                self.hass, timedelta(days=1), self._async_update
            )

    async def async_release_notes(self) -> str | None:
        """Get release notes."""
        if self._latest_version_firmware is None:
            return None
        return self._latest_version_firmware.changelog

    async def async_install(
        self, version: str | None, backup: bool, **kwargs: Any
    ) -> None:
        """Install an update."""
        firmware = self._latest_version_firmware
        assert firmware
        self._unsub_firmware_events_and_reset_progress(True)

        self._progress_unsub = self.node.on(
            "firmware update progress", self._update_progress
        )
        self._finished_unsub = self.node.once(
            "firmware update finished", self._update_finished
        )

        for file in firmware.files:
            try:
                await self.driver.controller.async_begin_ota_firmware_update(
                    self.node, file
                )
            except BaseZwaveJSServerError as err:
                self._unsub_firmware_events_and_reset_progress()
                raise HomeAssistantError(err) from err

            # We need to block until we receive the `firmware update finished` event
            await self._finished_event.wait()
            assert self._finished_status is not None

            # If status is not OK, we should throw an error to let the user know
            if self._finished_status not in (
                FirmwareUpdateStatus.OK_NO_RESTART,
                FirmwareUpdateStatus.OK_RESTART_PENDING,
                FirmwareUpdateStatus.OK_WAITING_FOR_ACTIVATION,
            ):
                status = self._finished_status
                self._unsub_firmware_events_and_reset_progress()
                raise HomeAssistantError(status.name.replace("_", " ").title())

            # If we get here, the firmware installation was successful and we need to
            # update progress accordingly
            self._num_files_installed += 1
            self._attr_in_progress = floor(
                100 * self._num_files_installed / len(firmware.files)
            )
            self.async_write_ha_state()

        self._attr_installed_version = self._attr_latest_version = firmware.version
        self._latest_version_firmware = None
        self._unsub_firmware_events_and_reset_progress()

    async def async_poll_value(self, _: bool) -> None:
        """Poll a value."""
        LOGGER.error(
            "There is no value to refresh for this entity so the zwave_js.refresh_value "
            "service won't work for it"
        )

    async def async_added_to_hass(self) -> None:
        """Call when entity is added."""
        self.async_on_remove(
            async_dispatcher_connect(
                self.hass,
                f"{DOMAIN}_{self.unique_id}_poll_value",
                self.async_poll_value,
            )
        )

        self.async_on_remove(
            async_dispatcher_connect(
                self.hass,
                f"{DOMAIN}_{self._base_unique_id}_remove_entity",
                self.async_remove,
            )
        )

        self.async_on_remove(
            async_dispatcher_connect(
                self.hass,
                f"{DOMAIN}_{self._base_unique_id}_remove_entity_on_ready_node",
                self.async_remove,
            )
        )

        self.async_on_remove(async_at_start(self.hass, self._async_update))

    async def async_will_remove_from_hass(self) -> None:
        """Call when entity will be removed."""
        if self._status_unsub:
            self._status_unsub()
            self._status_unsub = None

        if self._poll_unsub:
            self._poll_unsub()
            self._poll_unsub = None

        self._unsub_firmware_events_and_reset_progress()