Add esphome native device update entities (#119339)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/esphome/entry_data.py

@@ -38,6 +38,7 @@ from aioesphomeapi import (
     TextInfo,
     TextSensorInfo,
     TimeInfo,
+    UpdateInfo,
     UserService,
     ValveInfo,
     build_unique_id,
@@ -82,6 +83,7 @@ INFO_TYPE_TO_PLATFORM: dict[type[EntityInfo], Platform] = {
     TextInfo: Platform.TEXT,
     TextSensorInfo: Platform.SENSOR,
     TimeInfo: Platform.TIME,
+    UpdateInfo: Platform.UPDATE,
     ValveInfo: Platform.VALVE,
 }
 

homeassistant/components/esphome/update.py

@@ -5,7 +5,12 @@ from __future__ import annotations
 import asyncio
 from typing import Any
 
-from aioesphomeapi import DeviceInfo as ESPHomeDeviceInfo, EntityInfo
+from aioesphomeapi import (
+    DeviceInfo as ESPHomeDeviceInfo,
+    EntityInfo,
+    UpdateInfo,
+    UpdateState,
+)
 
 from homeassistant.components.update import (
     UpdateDeviceClass,
@@ -19,10 +24,17 @@ from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.util.enum import try_parse_enum
 
 from .coordinator import ESPHomeDashboardCoordinator
 from .dashboard import async_get_dashboard
 from .domain_data import DomainData
+from .entity import (
+    EsphomeEntity,
+    convert_api_error_ha_error,
+    esphome_state_property,
+    platform_async_setup_entry,
+)
 from .entry_data import RuntimeEntryData
 
 KEY_UPDATE_LOCK = "esphome_update_lock"
@@ -36,6 +48,15 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up ESPHome update based on a config entry."""
+    await platform_async_setup_entry(
+        hass,
+        entry,
+        async_add_entities,
+        info_type=UpdateInfo,
+        entity_type=ESPHomeUpdateEntity,
+        state_type=UpdateState,
+    )
+
     if (dashboard := async_get_dashboard(hass)) is None:
         return
     entry_data = DomainData.get(hass).get_entry_data(entry)
@@ -54,7 +75,7 @@ async def async_setup_entry(
             unsub()
         unsubs.clear()
 
-        async_add_entities([ESPHomeUpdateEntity(entry_data, dashboard)])
+        async_add_entities([ESPHomeDashboardUpdateEntity(entry_data, dashboard)])
 
     if entry_data.available and dashboard.last_update_success:
         _async_setup_update_entity()
@@ -66,7 +87,9 @@ async def async_setup_entry(
     ]
 
 
-class ESPHomeUpdateEntity(CoordinatorEntity[ESPHomeDashboardCoordinator], UpdateEntity):
+class ESPHomeDashboardUpdateEntity(
+    CoordinatorEntity[ESPHomeDashboardCoordinator], UpdateEntity
+):
     """Defines an ESPHome update entity."""
 
     _attr_has_entity_name = True
@@ -179,3 +202,65 @@ class ESPHomeUpdateEntity(CoordinatorEntity[ESPHomeDashboardCoordinator], Update
                     )
             finally:
                 await self.coordinator.async_request_refresh()
+
+
+class ESPHomeUpdateEntity(EsphomeEntity[UpdateInfo, UpdateState], UpdateEntity):
+    """A update implementation for esphome."""
+
+    _attr_supported_features = (
+        UpdateEntityFeature.INSTALL | UpdateEntityFeature.PROGRESS
+    )
+
+    @callback
+    def _on_static_info_update(self, static_info: EntityInfo) -> None:
+        """Set attrs from static info."""
+        super()._on_static_info_update(static_info)
+        static_info = self._static_info
+        self._attr_device_class = try_parse_enum(
+            UpdateDeviceClass, static_info.device_class
+        )
+
+    @property
+    @esphome_state_property
+    def installed_version(self) -> str | None:
+        """Return the installed version."""
+        return self._state.current_version
+
+    @property
+    @esphome_state_property
+    def in_progress(self) -> bool | int | None:
+        """Return if the update is in progress."""
+        if self._state.has_progress:
+            return int(self._state.progress)
+        return self._state.in_progress
+
+    @property
+    @esphome_state_property
+    def latest_version(self) -> str | None:
+        """Return the latest version."""
+        return self._state.latest_version
+
+    @property
+    @esphome_state_property
+    def release_summary(self) -> str | None:
+        """Return the release summary."""
+        return self._state.release_summary
+
+    @property
+    @esphome_state_property
+    def release_url(self) -> str | None:
+        """Return the release URL."""
+        return self._state.release_url
+
+    @property
+    @esphome_state_property
+    def title(self) -> str | None:
+        """Return the title of the update."""
+        return self._state.title
+
+    @convert_api_error_ha_error
+    async def async_install(
+        self, version: str | None, backup: bool, **kwargs: Any
+    ) -> None:
+        """Update the current value."""
+        self._client.update_command(key=self._key, install=True)

tests/components/esphome/test_update.py

@@ -3,12 +3,21 @@
 from collections.abc import Awaitable, Callable
 from unittest.mock import Mock, patch
 
-from aioesphomeapi import APIClient, EntityInfo, EntityState, UserService
+from aioesphomeapi import (
+    APIClient,
+    EntityInfo,
+    EntityState,
+    UpdateInfo,
+    UpdateState,
+    UserService,
+)
 import pytest
 
 from homeassistant.components.esphome.dashboard import async_get_dashboard
-from homeassistant.components.update import UpdateEntityFeature
+from homeassistant.components.update import DOMAIN as UPDATE_DOMAIN, UpdateEntityFeature
+from homeassistant.components.update.const import SERVICE_INSTALL
 from homeassistant.const import (
+    ATTR_ENTITY_ID,
     ATTR_SUPPORTED_FEATURES,
     STATE_OFF,
     STATE_ON,
@@ -83,7 +92,7 @@ async def test_update_entity(
 
     with patch(
         "homeassistant.components.esphome.update.DomainData.get_entry_data",
-        return_value=Mock(available=True, device_info=mock_device_info),
+        return_value=Mock(available=True, device_info=mock_device_info, info={}),
     ):
         assert await hass.config_entries.async_setup(mock_config_entry.entry_id)
         await hass.async_block_till_done()
@@ -266,7 +275,7 @@ async def test_update_entity_dashboard_not_available_startup(
     with (
         patch(
             "homeassistant.components.esphome.update.DomainData.get_entry_data",
-            return_value=Mock(available=True, device_info=mock_device_info),
+            return_value=Mock(available=True, device_info=mock_device_info, info={}),
         ),
         patch(
             "esphome_dashboard_api.ESPHomeDashboardAPI.get_devices",
@@ -358,7 +367,7 @@ async def test_update_entity_not_present_without_dashboard(
     """Test ESPHome update entity does not get created if there is no dashboard."""
     with patch(
         "homeassistant.components.esphome.update.DomainData.get_entry_data",
-        return_value=Mock(available=True, device_info=mock_device_info),
+        return_value=Mock(available=True, device_info=mock_device_info, info={}),
     ):
         assert await hass.config_entries.async_setup(mock_config_entry.entry_id)
         await hass.async_block_till_done()
@@ -408,3 +417,104 @@ async def test_update_becomes_available_at_runtime(
     # We now know the version so install is enabled
     features = state.attributes[ATTR_SUPPORTED_FEATURES]
     assert features is UpdateEntityFeature.INSTALL
+
+
+async def test_generic_device_update_entity(
+    hass: HomeAssistant,
+    mock_client: APIClient,
+    mock_generic_device_entry,
+) -> None:
+    """Test a generic device update entity."""
+    entity_info = [
+        UpdateInfo(
+            object_id="myupdate",
+            key=1,
+            name="my update",
+            unique_id="my_update",
+        )
+    ]
+    states = [
+        UpdateState(
+            key=1,
+            current_version="2024.6.0",
+            latest_version="2024.6.0",
+            title="ESPHome Project",
+            release_summary="This is a release summary",
+            release_url="https://esphome.io/changelog",
+        )
+    ]
+    user_service = []
+    await mock_generic_device_entry(
+        mock_client=mock_client,
+        entity_info=entity_info,
+        user_service=user_service,
+        states=states,
+    )
+    state = hass.states.get("update.test_myupdate")
+    assert state is not None
+    assert state.state == STATE_OFF
+
+
+async def test_generic_device_update_entity_has_update(
+    hass: HomeAssistant,
+    mock_client: APIClient,
+    mock_esphome_device: Callable[
+        [APIClient, list[EntityInfo], list[UserService], list[EntityState]],
+        Awaitable[MockESPHomeDevice],
+    ],
+) -> None:
+    """Test a generic device update entity with an update."""
+    entity_info = [
+        UpdateInfo(
+            object_id="myupdate",
+            key=1,
+            name="my update",
+            unique_id="my_update",
+        )
+    ]
+    states = [
+        UpdateState(
+            key=1,
+            current_version="2024.6.0",
+            latest_version="2024.6.1",
+            title="ESPHome Project",
+            release_summary="This is a release summary",
+            release_url="https://esphome.io/changelog",
+        )
+    ]
+    user_service = []
+    mock_device: MockESPHomeDevice = await mock_esphome_device(
+        mock_client=mock_client,
+        entity_info=entity_info,
+        user_service=user_service,
+        states=states,
+    )
+    state = hass.states.get("update.test_myupdate")
+    assert state is not None
+    assert state.state == STATE_ON
+
+    await hass.services.async_call(
+        UPDATE_DOMAIN,
+        SERVICE_INSTALL,
+        {ATTR_ENTITY_ID: "update.test_myupdate"},
+        blocking=True,
+    )
+
+    mock_device.set_state(
+        UpdateState(
+            key=1,
+            in_progress=True,
+            has_progress=True,
+            progress=50,
+            current_version="2024.6.0",
+            latest_version="2024.6.1",
+            title="ESPHome Project",
+            release_summary="This is a release summary",
+            release_url="https://esphome.io/changelog",
+        )
+    )
+
+    state = hass.states.get("update.test_myupdate")
+    assert state is not None
+    assert state.state == STATE_ON
+    assert state.attributes["in_progress"] == 50