Add streaming to Teslemetry update platform (#140021)

* Update platform

* Tests

* fix tests

homeassistant/components/teslemetry/update.py

@@ -2,16 +2,22 @@
 
 from __future__ import annotations
 
-from typing import Any, cast
+from typing import Any
 
 from tesla_fleet_api.const import Scope
+from tesla_fleet_api.vehiclespecific import VehicleSpecific
 
 from homeassistant.components.update import UpdateEntity, UpdateEntityFeature
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+from homeassistant.helpers.restore_state import RestoreEntity
 
 from . import TeslemetryConfigEntry
-from .entity import TeslemetryVehicleEntity
+from .entity import (
+    TeslemetryRootEntity,
+    TeslemetryVehicleEntity,
+    TeslemetryVehicleStreamEntity,
+)
 from .helpers import handle_vehicle_command
 from .models import TeslemetryVehicleData
 
@@ -32,12 +38,31 @@ async def async_setup_entry(
     """Set up the Teslemetry update platform from a config entry."""
 
     async_add_entities(
-        TeslemetryUpdateEntity(vehicle, entry.runtime_data.scopes)
+        TeslemetryPollingUpdateEntity(vehicle, entry.runtime_data.scopes)
+        if vehicle.api.pre2021 or vehicle.firmware < "2024.44.25"
+        else TeslemetryStreamingUpdateEntity(vehicle, entry.runtime_data.scopes)
         for vehicle in entry.runtime_data.vehicles
     )
 
 
-class TeslemetryUpdateEntity(TeslemetryVehicleEntity, UpdateEntity):
+class TeslemetryUpdateEntity(TeslemetryRootEntity, UpdateEntity):
+    """Teslemetry Updates entity."""
+
+    api: VehicleSpecific
+    _attr_supported_features = UpdateEntityFeature.PROGRESS
+
+    async def async_install(
+        self, version: str | None, backup: bool, **kwargs: Any
+    ) -> None:
+        """Install an update."""
+        self.raise_for_scope(Scope.VEHICLE_CMDS)
+
+        await handle_vehicle_command(self.api.schedule_software_update(offset_sec=0))
+        self._attr_in_progress = True
+        self.async_write_ha_state()
+
+
+class TeslemetryPollingUpdateEntity(TeslemetryVehicleEntity, TeslemetryUpdateEntity):
     """Teslemetry Updates entity."""
 
     def __init__(
@@ -94,18 +119,125 @@ class TeslemetryUpdateEntity(TeslemetryVehicleEntity, UpdateEntity):
         ):
             self._attr_in_progress = True
             if install_perc := self.get("vehicle_state_software_update_install_perc"):
-                self._attr_update_percentage = cast(int, install_perc)
+                self._attr_update_percentage = install_perc
         else:
             self._attr_in_progress = False
             self._attr_update_percentage = None
 
-    async def async_install(
-        self, version: str | None, backup: bool, **kwargs: Any
+
+class TeslemetryStreamingUpdateEntity(
+    TeslemetryVehicleStreamEntity, TeslemetryUpdateEntity, RestoreEntity
+):
+    """Teslemetry Updates entity."""
+
+    _download_percentage: int = 0
+    _install_percentage: int = 0
+
+    def __init__(
+        self,
+        data: TeslemetryVehicleData,
+        scopes: list[Scope],
     ) -> None:
-        """Install an update."""
-        self.raise_for_scope(Scope.ENERGY_CMDS)
-        await self.wake_up_if_asleep()
-        await handle_vehicle_command(self.api.schedule_software_update(offset_sec=60))
-        self._attr_in_progress = True
-        self._attr_update_percentage = None
+        """Initialize the Update."""
+        self.scoped = Scope.VEHICLE_CMDS in scopes
+        super().__init__(
+            data,
+            "vehicle_state_software_update_status",
+        )
+
+    async def async_added_to_hass(self) -> None:
+        """Handle entity which will be added."""
+        await super().async_added_to_hass()
+        if (state := await self.async_get_last_state()) is not None:
+            self._attr_in_progress = state.attributes.get("in_progress", False)
+            self._install_percentage = state.attributes.get("install_percentage", False)
+            self._attr_installed_version = state.attributes.get("installed_version")
+            self._attr_latest_version = state.attributes.get("latest_version")
+            self._attr_supported_features = UpdateEntityFeature(
+                state.attributes.get(
+                    "supported_features", self._attr_supported_features
+                )
+            )
             self.async_write_ha_state()
+
+        self.async_on_remove(
+            self.vehicle.stream_vehicle.listen_SoftwareUpdateDownloadPercentComplete(
+                self._async_handle_software_update_download_percent_complete
+            )
+        )
+        self.async_on_remove(
+            self.vehicle.stream_vehicle.listen_SoftwareUpdateInstallationPercentComplete(
+                self._async_handle_software_update_installation_percent_complete
+            )
+        )
+        self.async_on_remove(
+            self.vehicle.stream_vehicle.listen_SoftwareUpdateScheduledStartTime(
+                self._async_handle_software_update_scheduled_start_time
+            )
+        )
+        self.async_on_remove(
+            self.vehicle.stream_vehicle.listen_SoftwareUpdateVersion(
+                self._async_handle_software_update_version
+            )
+        )
+        self.async_on_remove(
+            self.vehicle.stream_vehicle.listen_Version(self._async_handle_version)
+        )
+
+    def _async_handle_software_update_download_percent_complete(
+        self, value: float | None
+    ):
+        """Handle software update download percent complete."""
+
+        self._download_percentage = round(value) if value is not None else 0
+        if self.scoped and self._download_percentage == 100:
+            self._attr_supported_features = (
+                UpdateEntityFeature.PROGRESS | UpdateEntityFeature.INSTALL
+            )
+        else:
+            self._attr_supported_features = UpdateEntityFeature.PROGRESS
+        self._async_update_progress()
+        self.async_write_ha_state()
+
+    def _async_handle_software_update_installation_percent_complete(
+        self, value: float | None
+    ):
+        """Handle software update installation percent complete."""
+
+        self._install_percentage = round(value) if value is not None else 0
+        self._async_update_progress()
+        self.async_write_ha_state()
+
+    def _async_handle_software_update_scheduled_start_time(self, value: str | None):
+        """Handle software update scheduled start time."""
+
+        self._attr_in_progress = value is not None
+        self.async_write_ha_state()
+
+    def _async_handle_software_update_version(self, value: str | None):
+        """Handle software update version."""
+
+        self._attr_latest_version = (
+            value if value and value != " " else self._attr_installed_version
+        )
+        self.async_write_ha_state()
+
+    def _async_handle_version(self, value: str | None):
+        """Handle version."""
+
+        if value is not None:
+            self._attr_installed_version = value.split(" ")[0]
+            self.async_write_ha_state()
+
+    def _async_update_progress(self) -> None:
+        """Update the progress of the update."""
+
+        if self._download_percentage > 1 and self._download_percentage < 100:
+            self._attr_in_progress = True
+            self._attr_update_percentage = self._download_percentage
+        elif self._install_percentage > 1:
+            self._attr_in_progress = True
+            self._attr_update_percentage = self._install_percentage
+        else:
+            self._attr_in_progress = False
+            self._attr_update_percentage = None

tests/components/teslemetry/snapshots/test_update.ambr

@@ -117,3 +117,128 @@
     'state': 'off',
   })
 # ---
+# name: test_update_streaming[downloading]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'auto_update': False,
+      'display_precision': 0,
+      'entity_picture': 'https://brands.home-assistant.io/_/teslemetry/icon.png',
+      'friendly_name': 'Test Update',
+      'in_progress': False,
+      'installed_version': '2025.1.1',
+      'latest_version': '2025.2.1',
+      'release_summary': None,
+      'release_url': None,
+      'skipped_version': None,
+      'supported_features': <UpdateEntityFeature: 4>,
+      'title': None,
+      'update_percentage': None,
+    }),
+    'context': <ANY>,
+    'entity_id': 'update.test_update',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'on',
+  })
+# ---
+# name: test_update_streaming[installing]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'auto_update': False,
+      'display_precision': 0,
+      'entity_picture': 'https://brands.home-assistant.io/_/teslemetry/icon.png',
+      'friendly_name': 'Test Update',
+      'in_progress': False,
+      'installed_version': '2025.1.1',
+      'latest_version': '2025.2.1',
+      'release_summary': None,
+      'release_url': None,
+      'skipped_version': None,
+      'supported_features': <UpdateEntityFeature: 5>,
+      'title': None,
+      'update_percentage': None,
+    }),
+    'context': <ANY>,
+    'entity_id': 'update.test_update',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'on',
+  })
+# ---
+# name: test_update_streaming[ready]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'auto_update': False,
+      'display_precision': 0,
+      'entity_picture': 'https://brands.home-assistant.io/_/teslemetry/icon.png',
+      'friendly_name': 'Test Update',
+      'in_progress': False,
+      'installed_version': '2025.1.1',
+      'latest_version': '2025.2.1',
+      'release_summary': None,
+      'release_url': None,
+      'skipped_version': None,
+      'supported_features': <UpdateEntityFeature: 5>,
+      'title': None,
+      'update_percentage': None,
+    }),
+    'context': <ANY>,
+    'entity_id': 'update.test_update',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'on',
+  })
+# ---
+# name: test_update_streaming[restored]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'auto_update': False,
+      'display_precision': 0,
+      'entity_picture': 'https://brands.home-assistant.io/_/teslemetry/icon.png',
+      'friendly_name': 'Test Update',
+      'in_progress': False,
+      'installed_version': '2025.2.1',
+      'latest_version': '2025.1.1',
+      'release_summary': None,
+      'release_url': None,
+      'skipped_version': None,
+      'supported_features': <UpdateEntityFeature: 4>,
+      'title': None,
+      'update_percentage': None,
+    }),
+    'context': <ANY>,
+    'entity_id': 'update.test_update',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'off',
+  })
+# ---
+# name: test_update_streaming[updated]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'auto_update': False,
+      'display_precision': 0,
+      'entity_picture': 'https://brands.home-assistant.io/_/teslemetry/icon.png',
+      'friendly_name': 'Test Update',
+      'in_progress': False,
+      'installed_version': '2025.2.1',
+      'latest_version': '2025.1.1',
+      'release_summary': None,
+      'release_url': None,
+      'skipped_version': None,
+      'supported_features': <UpdateEntityFeature: 4>,
+      'title': None,
+      'update_percentage': None,
+    }),
+    'context': <ANY>,
+    'entity_id': 'update.test_update',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'off',
+  })
+# ---

tests/components/teslemetry/test_update.py

@@ -4,7 +4,9 @@ import copy
 from unittest.mock import AsyncMock, patch
 
 from freezegun.api import FrozenDateTimeFactory
+import pytest
 from syrupy.assertion import SnapshotAssertion
+from teslemetry_stream import Signal
 
 from homeassistant.components.teslemetry.coordinator import VEHICLE_INTERVAL
 from homeassistant.components.teslemetry.update import INSTALLING
@@ -13,7 +15,7 @@ from homeassistant.const import ATTR_ENTITY_ID, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 
-from . import assert_entities, setup_platform
+from . import assert_entities, reload_platform, setup_platform
 from .const import COMMAND_OK, VEHICLE_DATA, VEHICLE_DATA_ALT
 
 from tests.common import async_fire_time_changed
@@ -23,6 +25,7 @@ async def test_update(
     hass: HomeAssistant,
     snapshot: SnapshotAssertion,
     entity_registry: er.EntityRegistry,
+    mock_legacy: AsyncMock,
 ) -> None:
     """Tests that the update entities are correct."""
 
@@ -35,6 +38,7 @@ async def test_update_alt(
     snapshot: SnapshotAssertion,
     entity_registry: er.EntityRegistry,
     mock_vehicle_data: AsyncMock,
+    mock_legacy: AsyncMock,
 ) -> None:
     """Tests that the update entities are correct."""
 
@@ -48,6 +52,7 @@ async def test_update_services(
     mock_vehicle_data: AsyncMock,
     freezer: FrozenDateTimeFactory,
     snapshot: SnapshotAssertion,
+    mock_legacy: AsyncMock,
 ) -> None:
     """Tests that the update services work."""
 
@@ -78,3 +83,90 @@ async def test_update_services(
 
     state = hass.states.get(entity_id)
     assert state.attributes["in_progress"] == 1
+
+
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+async def test_update_streaming(
+    hass: HomeAssistant,
+    snapshot: SnapshotAssertion,
+    mock_vehicle_data: AsyncMock,
+    mock_add_listener: AsyncMock,
+) -> None:
+    """Tests that the select entities with streaming are correct."""
+
+    entry = await setup_platform(hass, [Platform.UPDATE])
+
+    # Stream update
+    mock_add_listener.send(
+        {
+            "vin": VEHICLE_DATA_ALT["response"]["vin"],
+            "data": {
+                Signal.SOFTWARE_UPDATE_DOWNLOAD_PERCENT_COMPLETE: 50,
+                Signal.SOFTWARE_UPDATE_INSTALLATION_PERCENT_COMPLETE: None,
+                Signal.SOFTWARE_UPDATE_SCHEDULED_START_TIME: None,
+                Signal.SOFTWARE_UPDATE_VERSION: "2025.2.1",
+                Signal.VERSION: "2025.1.1",
+            },
+            "createdAt": "2024-10-04T10:45:17.537Z",
+        }
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("update.test_update")
+    assert state == snapshot(name="downloading")
+
+    mock_add_listener.send(
+        {
+            "vin": VEHICLE_DATA_ALT["response"]["vin"],
+            "data": {
+                Signal.SOFTWARE_UPDATE_DOWNLOAD_PERCENT_COMPLETE: 100,
+                Signal.SOFTWARE_UPDATE_INSTALLATION_PERCENT_COMPLETE: 1,
+                Signal.SOFTWARE_UPDATE_SCHEDULED_START_TIME: None,
+                Signal.SOFTWARE_UPDATE_VERSION: "2025.2.1",
+                Signal.VERSION: "2025.1.1",
+            },
+            "createdAt": "2024-10-04T10:45:17.537Z",
+        }
+    )
+    await hass.async_block_till_done()
+    state = hass.states.get("update.test_update")
+    assert state == snapshot(name="ready")
+
+    mock_add_listener.send(
+        {
+            "vin": VEHICLE_DATA_ALT["response"]["vin"],
+            "data": {
+                Signal.SOFTWARE_UPDATE_DOWNLOAD_PERCENT_COMPLETE: 100,
+                Signal.SOFTWARE_UPDATE_INSTALLATION_PERCENT_COMPLETE: 50,
+                Signal.SOFTWARE_UPDATE_SCHEDULED_START_TIME: None,
+                Signal.SOFTWARE_UPDATE_VERSION: "2025.2.1",
+                Signal.VERSION: "2025.1.1",
+            },
+            "createdAt": "2024-10-04T10:45:17.537Z",
+        }
+    )
+    await hass.async_block_till_done()
+    state = hass.states.get("update.test_update")
+    assert state == snapshot(name="installing")
+
+    mock_add_listener.send(
+        {
+            "vin": VEHICLE_DATA_ALT["response"]["vin"],
+            "data": {
+                Signal.SOFTWARE_UPDATE_DOWNLOAD_PERCENT_COMPLETE: None,
+                Signal.SOFTWARE_UPDATE_INSTALLATION_PERCENT_COMPLETE: None,
+                Signal.SOFTWARE_UPDATE_SCHEDULED_START_TIME: None,
+                Signal.SOFTWARE_UPDATE_VERSION: "",
+                Signal.VERSION: "2025.2.1",
+            },
+            "createdAt": "2024-10-04T10:45:17.537Z",
+        }
+    )
+    await hass.async_block_till_done()
+    state = hass.states.get("update.test_update")
+    assert state == snapshot(name="updated")
+
+    await reload_platform(hass, entry, [Platform.UPDATE])
+
+    state = hass.states.get("update.test_update")
+    assert state == snapshot(name="restored")