Add streaming select entities to Teslemetry (#137210)

homeassistant/components/teslemetry/select.py

@@ -2,18 +2,27 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from itertools import chain
+from typing import Any
 
+from tesla_fleet_api import VehicleSpecific
 from tesla_fleet_api.const import EnergyExportMode, EnergyOperationMode, Scope, Seat
+from teslemetry_stream import TeslemetryStreamVehicle
 
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.restore_state import RestoreEntity
 
 from . import TeslemetryConfigEntry
-from .entity import TeslemetryEnergyInfoEntity, TeslemetryVehicleEntity
+from .entity import (
+    TeslemetryEnergyInfoEntity,
+    TeslemetryRootEntity,
+    TeslemetryVehicleEntity,
+    TeslemetryVehicleStreamEntity,
+)
 from .helpers import handle_command, handle_vehicle_command
 from .models import TeslemetryEnergyData, TeslemetryVehicleData
 
@@ -24,53 +33,136 @@ HIGH = "high"
 
 PARALLEL_UPDATES = 0
 
+LEVEL = {OFF: 0, LOW: 1, MEDIUM: 2, HIGH: 3}
+
 
 @dataclass(frozen=True, kw_only=True)
-class SeatHeaterDescription(SelectEntityDescription):
+class TeslemetrySelectEntityDescription(SelectEntityDescription):
     """Seat Heater entity description."""
 
-    position: Seat
+    select_fn: Callable[[VehicleSpecific, int], Awaitable[Any]]
-    available_fn: Callable[[TeslemetrySeatHeaterSelectEntity], bool] = lambda _: True
+    supported_fn: Callable[[dict], bool] = lambda _: True
+    streaming_listener: (
+        Callable[
+            [TeslemetryStreamVehicle, Callable[[int | None], None]],
+            Callable[[], None],
+        ]
+        | None
+    ) = None
+    options: list[str]
 
 
-SEAT_HEATER_DESCRIPTIONS: tuple[SeatHeaterDescription, ...] = (
+VEHICLE_DESCRIPTIONS: tuple[TeslemetrySelectEntityDescription, ...] = (
-    SeatHeaterDescription(
+    TeslemetrySelectEntityDescription(
         key="climate_state_seat_heater_left",
-        position=Seat.FRONT_LEFT,
+        select_fn=lambda api, level: api.remote_seat_heater_request(
+            Seat.FRONT_LEFT, level
+        ),
+        streaming_listener=lambda x, y: x.listen_SeatHeaterLeft(y),
+        options=[
+            OFF,
+            LOW,
+            MEDIUM,
+            HIGH,
+        ],
     ),
-    SeatHeaterDescription(
+    TeslemetrySelectEntityDescription(
         key="climate_state_seat_heater_right",
-        position=Seat.FRONT_RIGHT,
+        select_fn=lambda api, level: api.remote_seat_heater_request(
+            Seat.FRONT_RIGHT, level
+        ),
+        streaming_listener=lambda x, y: x.listen_SeatHeaterRight(y),
+        options=[
+            OFF,
+            LOW,
+            MEDIUM,
+            HIGH,
+        ],
     ),
-    SeatHeaterDescription(
+    TeslemetrySelectEntityDescription(
         key="climate_state_seat_heater_rear_left",
-        position=Seat.REAR_LEFT,
+        select_fn=lambda api, level: api.remote_seat_heater_request(
-        available_fn=lambda self: self.get("vehicle_config_rear_seat_heaters") != 0,
+            Seat.REAR_LEFT, level
+        ),
+        supported_fn=lambda data: data.get("vehicle_config_rear_seat_heaters") != 0,
+        streaming_listener=lambda x, y: x.listen_SeatHeaterRearLeft(y),
         entity_registry_enabled_default=False,
+        options=[
+            OFF,
+            LOW,
+            MEDIUM,
+            HIGH,
+        ],
     ),
-    SeatHeaterDescription(
+    TeslemetrySelectEntityDescription(
         key="climate_state_seat_heater_rear_center",
-        position=Seat.REAR_CENTER,
+        select_fn=lambda api, level: api.remote_seat_heater_request(
-        available_fn=lambda self: self.get("vehicle_config_rear_seat_heaters") != 0,
+            Seat.REAR_CENTER, level
+        ),
+        supported_fn=lambda data: data.get("vehicle_config_rear_seat_heaters") != 0,
+        streaming_listener=lambda x, y: x.listen_SeatHeaterRearCenter(y),
         entity_registry_enabled_default=False,
+        options=[
+            OFF,
+            LOW,
+            MEDIUM,
+            HIGH,
+        ],
     ),
-    SeatHeaterDescription(
+    TeslemetrySelectEntityDescription(
         key="climate_state_seat_heater_rear_right",
-        position=Seat.REAR_RIGHT,
+        select_fn=lambda api, level: api.remote_seat_heater_request(
-        available_fn=lambda self: self.get("vehicle_config_rear_seat_heaters") != 0,
+            Seat.REAR_RIGHT, level
+        ),
+        supported_fn=lambda data: data.get("vehicle_config_rear_seat_heaters") != 0,
+        streaming_listener=lambda x, y: x.listen_SeatHeaterRearRight(y),
         entity_registry_enabled_default=False,
+        options=[
+            OFF,
+            LOW,
+            MEDIUM,
+            HIGH,
+        ],
     ),
-    SeatHeaterDescription(
+    TeslemetrySelectEntityDescription(
         key="climate_state_seat_heater_third_row_left",
-        position=Seat.THIRD_LEFT,
+        select_fn=lambda api, level: api.remote_seat_heater_request(
-        available_fn=lambda self: self.get("vehicle_config_third_row_seats") != "None",
+            Seat.THIRD_LEFT, level
+        ),
+        supported_fn=lambda self: self.get("vehicle_config_third_row_seats") != "None",
         entity_registry_enabled_default=False,
+        options=[
+            OFF,
+            LOW,
+            MEDIUM,
+            HIGH,
+        ],
     ),
-    SeatHeaterDescription(
+    TeslemetrySelectEntityDescription(
         key="climate_state_seat_heater_third_row_right",
-        position=Seat.THIRD_RIGHT,
+        select_fn=lambda api, level: api.remote_seat_heater_request(
-        available_fn=lambda self: self.get("vehicle_config_third_row_seats") != "None",
+            Seat.THIRD_RIGHT, level
+        ),
+        supported_fn=lambda self: self.get("vehicle_config_third_row_seats") != "None",
         entity_registry_enabled_default=False,
+        options=[
+            OFF,
+            LOW,
+            MEDIUM,
+            HIGH,
+        ],
+    ),
+    TeslemetrySelectEntityDescription(
+        key="climate_state_steering_wheel_heat_level",
+        select_fn=lambda api, level: api.remote_steering_wheel_heat_level_request(
+            level
+        ),
+        streaming_listener=lambda x, y: x.listen_HvacSteeringWheelHeatLevel(y),
+        options=[
+            OFF,
+            LOW,
+            HIGH,
+        ],
     ),
 )
 
@@ -85,17 +177,18 @@ async def async_setup_entry(
     async_add_entities(
         chain(
             (
-                TeslemetrySeatHeaterSelectEntity(
+                TeslemetryPollingSelectEntity(
                     vehicle, description, entry.runtime_data.scopes
                 )
-                for description in SEAT_HEATER_DESCRIPTIONS
+                if vehicle.api.pre2021
+                or vehicle.firmware < "2024.26"
+                or description.streaming_listener is None
+                else TeslemetryStreamingSelectEntity(
+                    vehicle, description, entry.runtime_data.scopes
+                )
+                for description in VEHICLE_DESCRIPTIONS
                 for vehicle in entry.runtime_data.vehicles
-                if description.key in vehicle.coordinator.data
+                if description.supported_fn(vehicle.coordinator.data)
-            ),
-            (
-                TeslemetryWheelHeaterSelectEntity(vehicle, entry.runtime_data.scopes)
-                for vehicle in entry.runtime_data.vehicles
-                if vehicle.coordinator.data.get("climate_state_steering_wheel_heater")
             ),
             (
                 TeslemetryOperationSelectEntity(energysite, entry.runtime_data.scopes)
@@ -112,22 +205,31 @@ async def async_setup_entry(
     )
 
 
-class TeslemetrySeatHeaterSelectEntity(TeslemetryVehicleEntity, SelectEntity):
+class TeslemetrySelectEntity(TeslemetryRootEntity, SelectEntity):
-    """Select entity for vehicle seat heater."""
+    """Parent vehicle select entity class."""
 
-    entity_description: SeatHeaterDescription
+    entity_description: TeslemetrySelectEntityDescription
+    _climate: bool = False
 
-    _attr_options = [
+    async def async_select_option(self, option: str) -> None:
-        OFF,
+        """Change the selected option."""
-        LOW,
+        self.raise_for_scope(Scope.VEHICLE_CMDS)
-        MEDIUM,
+        level = LEVEL[option]
-        HIGH,
+        # AC must be on to turn on heaters
-    ]
+        if level and not self._climate:
+            await handle_vehicle_command(self.api.auto_conditioning_start())
+        await handle_vehicle_command(self.entity_description.select_fn(self.api, level))
+        self._attr_current_option = option
+        self.async_write_ha_state()
+
+
+class TeslemetryPollingSelectEntity(TeslemetryVehicleEntity, TeslemetrySelectEntity):
+    """Base polling vehicle select entity class."""
 
     def __init__(
         self,
         data: TeslemetryVehicleData,
-        description: SeatHeaterDescription,
+        description: TeslemetrySelectEntityDescription,
         scopes: list[Scope],
     ) -> None:
         """Initialize the vehicle seat select entity."""
@@ -137,72 +239,63 @@ class TeslemetrySeatHeaterSelectEntity(TeslemetryVehicleEntity, SelectEntity):
 
     def _async_update_attrs(self) -> None:
         """Handle updated data from the coordinator."""
-        self._attr_available = self.entity_description.available_fn(self)
+        self._climate = bool(self.get("climate_state_is_climate_on"))
-        value = self._value
+        if not isinstance(self._value, int):
-        if not isinstance(value, int):
             self._attr_current_option = None
         else:
-            self._attr_current_option = self._attr_options[value]
+            self._attr_current_option = self.entity_description.options[self._value]
-
-    async def async_select_option(self, option: str) -> None:
-        """Change the selected option."""
-        self.raise_for_scope(Scope.VEHICLE_CMDS)
-        await self.wake_up_if_asleep()
-        level = self._attr_options.index(option)
-        # AC must be on to turn on seat heater
-        if level and not self.get("climate_state_is_climate_on"):
-            await handle_vehicle_command(self.api.auto_conditioning_start())
-        await handle_vehicle_command(
-            self.api.remote_seat_heater_request(self.entity_description.position, level)
-        )
-        self._attr_current_option = option
-        self.async_write_ha_state()
 
 
-class TeslemetryWheelHeaterSelectEntity(TeslemetryVehicleEntity, SelectEntity):
+class TeslemetryStreamingSelectEntity(
-    """Select entity for vehicle steering wheel heater."""
+    TeslemetryVehicleStreamEntity, TeslemetrySelectEntity, RestoreEntity
-
+):
-    _attr_options = [
+    """Base streaming vehicle select entity class."""
-        OFF,
-        LOW,
-        HIGH,
-    ]
 
     def __init__(
         self,
         data: TeslemetryVehicleData,
+        description: TeslemetrySelectEntityDescription,
         scopes: list[Scope],
     ) -> None:
-        """Initialize the vehicle steering wheel select entity."""
+        """Initialize the vehicle seat select entity."""
+        self.entity_description = description
         self.scoped = Scope.VEHICLE_CMDS in scopes
-        super().__init__(
+        self._attr_current_option = None
-            data,
+        super().__init__(data, description.key)
-            "climate_state_steering_wheel_heat_level",
+
+    async def async_added_to_hass(self) -> None:
+        """Handle entity which will be added."""
+        await super().async_added_to_hass()
+
+        # Restore state
+        if (state := await self.async_get_last_state()) is not None:
+            if state.state in self.entity_description.options:
+                self._attr_current_option = state.state
+
+        # Listen for streaming data
+        assert self.entity_description.streaming_listener is not None
+        self.async_on_remove(
+            self.entity_description.streaming_listener(
+                self.vehicle.stream_vehicle, self._value_callback
+            )
         )
 
-    def _async_update_attrs(self) -> None:
+        self.async_on_remove(
-        """Handle updated data from the coordinator."""
+            self.vehicle.stream_vehicle.listen_HvacACEnabled(self._climate_callback)
+        )
 
-        value = self._value
+    def _value_callback(self, value: int | None) -> None:
-        if not isinstance(value, int):
+        """Update the value of the entity."""
+        if value is None:
             self._attr_current_option = None
         else:
-            self._attr_current_option = self._attr_options[value]
+            self._attr_current_option = self.entity_description.options[value]
-
-    async def async_select_option(self, option: str) -> None:
-        """Change the selected option."""
-        self.raise_for_scope(Scope.VEHICLE_CMDS)
-        await self.wake_up_if_asleep()
-        level = self._attr_options.index(option)
-        # AC must be on to turn on steering wheel heater
-        if level and not self.get("climate_state_is_climate_on"):
-            await handle_vehicle_command(self.api.auto_conditioning_start())
-        await handle_vehicle_command(
-            self.api.remote_steering_wheel_heat_level_request(level)
-        )
-        self._attr_current_option = option
         self.async_write_ha_state()
 
+    def _climate_callback(self, value: bool | None) -> None:
+        """Update the value of the entity."""
+        self._climate = bool(value)
+
 
 class TeslemetryOperationSelectEntity(TeslemetryEnergyInfoEntity, SelectEntity):
     """Select entity for operation mode select entities."""

tests/components/teslemetry/snapshots/test_select.ambr

@@ -408,3 +408,78 @@
     'state': 'off',
   })
 # ---
+# name: test_select[select.test_steering_wheel_heater-entry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': dict({
+      'options': list([
+        'off',
+        'low',
+        'high',
+      ]),
+    }),
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'select',
+    'entity_category': None,
+    'entity_id': 'select.test_steering_wheel_heater',
+    'has_entity_name': True,
+    'hidden_by': None,
+    'icon': None,
+    'id': <ANY>,
+    'labels': set({
+    }),
+    'name': None,
+    'options': dict({
+    }),
+    'original_device_class': None,
+    'original_icon': None,
+    'original_name': 'Steering wheel heater',
+    'platform': 'teslemetry',
+    'previous_unique_id': None,
+    'supported_features': 0,
+    'translation_key': 'climate_state_steering_wheel_heat_level',
+    'unique_id': 'LRW3F7EK4NC700000-climate_state_steering_wheel_heat_level',
+    'unit_of_measurement': None,
+  })
+# ---
+# name: test_select[select.test_steering_wheel_heater-state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Test Steering wheel heater',
+      'options': list([
+        'off',
+        'low',
+        'high',
+      ]),
+    }),
+    'context': <ANY>,
+    'entity_id': 'select.test_steering_wheel_heater',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'off',
+  })
+# ---
+# name: test_select_streaming[select.test_seat_heater_front_left]
+  'off'
+# ---
+# name: test_select_streaming[select.test_seat_heater_front_right]
+  'low'
+# ---
+# name: test_select_streaming[select.test_seat_heater_rear_center]
+  'unknown'
+# ---
+# name: test_select_streaming[select.test_seat_heater_rear_left]
+  'medium'
+# ---
+# name: test_select_streaming[select.test_seat_heater_rear_right]
+  'high'
+# ---
+# name: test_select_streaming[select.test_steering_wheel_heater]
+  'off'
+# ---

tests/components/teslemetry/test_select.py

@@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, patch
 import pytest
 from syrupy.assertion import SnapshotAssertion
 from tesla_fleet_api.const import EnergyExportMode, EnergyOperationMode
+from teslemetry_stream.const import Signal
 
 from homeassistant.components.select import (
     ATTR_OPTION,
@@ -16,7 +17,7 @@ from homeassistant.const import ATTR_ENTITY_ID, STATE_UNKNOWN, 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_ALT
 
 
@@ -25,6 +26,7 @@ async def test_select(
     hass: HomeAssistant,
     snapshot: SnapshotAssertion,
     entity_registry: er.EntityRegistry,
+    mock_legacy: AsyncMock,
 ) -> None:
     """Tests that the select entities are correct."""
 
@@ -106,6 +108,7 @@ async def test_select_invalid_data(
     snapshot: SnapshotAssertion,
     entity_registry: er.EntityRegistry,
     mock_vehicle_data: AsyncMock,
+    mock_legacy: AsyncMock,
 ) -> None:
     """Tests that the select entities handle invalid data."""
 
@@ -119,3 +122,45 @@ async def test_select_invalid_data(
     assert state.state == STATE_UNKNOWN
     state = hass.states.get("select.test_steering_wheel_heater")
     assert state.state == STATE_UNKNOWN
+
+
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+async def test_select_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.SELECT])
+
+    # Stream update
+    mock_add_listener.send(
+        {
+            "vin": VEHICLE_DATA_ALT["response"]["vin"],
+            "data": {
+                Signal.SEAT_HEATER_LEFT: 0,
+                Signal.SEAT_HEATER_RIGHT: 1,
+                Signal.SEAT_HEATER_REAR_LEFT: 2,
+                Signal.SEAT_HEATER_REAR_RIGHT: 3,
+                Signal.HVAC_STEERING_WHEEL_HEAT_LEVEL: 0,
+            },
+            "createdAt": "2024-10-04T10:45:17.537Z",
+        }
+    )
+    await hass.async_block_till_done()
+
+    await reload_platform(hass, entry, [Platform.SELECT])
+
+    # Assert the entities restored their values
+    for entity_id in (
+        "select.test_seat_heater_front_left",
+        "select.test_seat_heater_front_right",
+        "select.test_seat_heater_rear_left",
+        "select.test_seat_heater_rear_center",
+        "select.test_seat_heater_rear_right",
+        "select.test_steering_wheel_heater",
+    ):
+        state = hass.states.get(entity_id)
+        assert state.state == snapshot(name=entity_id)