Add streaming switches to Teslemetry (#137145)

* Add streaming switches

* Add switch tests

* Update snapshot

* Fix sentry

* update test docstring

homeassistant/components/teslemetry/switch.py

@@ -8,6 +8,7 @@ from itertools import chain
 from typing import Any
 
 from tesla_fleet_api.const import Scope, Seat
+from teslemetry_stream import TeslemetryStreamVehicle
 
 from homeassistant.components.switch import (
     SwitchDeviceClass,
@@ -16,10 +17,16 @@ from homeassistant.components.switch import (
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.typing import StateType
 
 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
 
@@ -34,18 +41,27 @@ class TeslemetrySwitchEntityDescription(SwitchEntityDescription):
     off_func: Callable
     scopes: list[Scope]
     value_func: Callable[[StateType], bool] = bool
+    streaming_listener: Callable[
+        [TeslemetryStreamVehicle, Callable[[StateType], None]],
+        Callable[[], None],
+    ]
+    streaming_value_fn: Callable[[StateType], bool] = bool
+    streaming_firmware: str = "2024.26"
     unique_id: str | None = None
 
 
 VEHICLE_DESCRIPTIONS: tuple[TeslemetrySwitchEntityDescription, ...] = (
     TeslemetrySwitchEntityDescription(
         key="vehicle_state_sentry_mode",
+        streaming_listener=lambda x, y: x.listen_SentryMode(y),
+        streaming_value_fn=lambda x: x != "Off",
         on_func=lambda api: api.set_sentry_mode(on=True),
         off_func=lambda api: api.set_sentry_mode(on=False),
         scopes=[Scope.VEHICLE_CMDS],
     ),
     TeslemetrySwitchEntityDescription(
         key="climate_state_auto_seat_climate_left",
+        streaming_listener=lambda x, y: x.listen_AutoSeatClimateLeft(y),
         on_func=lambda api: api.remote_auto_seat_climate_request(Seat.FRONT_LEFT, True),
         off_func=lambda api: api.remote_auto_seat_climate_request(
             Seat.FRONT_LEFT, False
@@ -54,6 +70,7 @@ VEHICLE_DESCRIPTIONS: tuple[TeslemetrySwitchEntityDescription, ...] = (
     ),
     TeslemetrySwitchEntityDescription(
         key="climate_state_auto_seat_climate_right",
+        streaming_listener=lambda x, y: x.listen_AutoSeatClimateRight(y),
         on_func=lambda api: api.remote_auto_seat_climate_request(
             Seat.FRONT_RIGHT, True
         ),
@@ -64,6 +81,7 @@ VEHICLE_DESCRIPTIONS: tuple[TeslemetrySwitchEntityDescription, ...] = (
     ),
     TeslemetrySwitchEntityDescription(
         key="climate_state_auto_steering_wheel_heat",
+        streaming_listener=lambda x, y: x.listen_HvacSteeringWheelHeatAuto(y),
         on_func=lambda api: api.remote_auto_steering_wheel_heat_climate_request(
             on=True
         ),
@@ -74,6 +92,8 @@ VEHICLE_DESCRIPTIONS: tuple[TeslemetrySwitchEntityDescription, ...] = (
     ),
     TeslemetrySwitchEntityDescription(
         key="climate_state_defrost_mode",
+        streaming_listener=lambda x, y: x.listen_DefrostMode(y),
+        streaming_value_fn=lambda x: x != "Off",
         on_func=lambda api: api.set_preconditioning_max(on=True, manual_override=False),
         off_func=lambda api: api.set_preconditioning_max(
             on=False, manual_override=False
@@ -83,9 +103,11 @@ VEHICLE_DESCRIPTIONS: tuple[TeslemetrySwitchEntityDescription, ...] = (
     TeslemetrySwitchEntityDescription(
         key="charge_state_charging_state",
         unique_id="charge_state_user_charge_enable_request",
+        value_func=lambda state: state in {"Starting", "Charging"},
+        streaming_listener=lambda x, y: x.listen_DetailedChargeState(y),
+        streaming_value_fn=lambda x: x in {"Starting", "Charging"},
         on_func=lambda api: api.charge_start(),
         off_func=lambda api: api.charge_stop(),
-        value_func=lambda state: state in {"Starting", "Charging"},
         scopes=[Scope.VEHICLE_CMDS, Scope.VEHICLE_CHARGING_CMDS],
     ),
 )
@@ -101,12 +123,16 @@ async def async_setup_entry(
     async_add_entities(
         chain(
             (
-                TeslemetryVehicleSwitchEntity(
+                TeslemetryPollingVehicleSwitchEntity(
+                    vehicle, description, entry.runtime_data.scopes
+                )
+                if vehicle.api.pre2021
+                or vehicle.firmware < description.streaming_firmware
+                else TeslemetryStreamingVehicleSwitchEntity(
                     vehicle, description, entry.runtime_data.scopes
                 )
                 for vehicle in entry.runtime_data.vehicles
                 for description in VEHICLE_DESCRIPTIONS
-                if description.key in vehicle.coordinator.data
             ),
             (
                 TeslemetryChargeFromGridSwitchEntity(
@@ -126,15 +152,31 @@ async def async_setup_entry(
     )
 
 
-class TeslemetrySwitchEntity(SwitchEntity):
+class TeslemetryVehicleSwitchEntity(TeslemetryRootEntity, SwitchEntity):
     """Base class for all Teslemetry switch entities."""
 
     _attr_device_class = SwitchDeviceClass.SWITCH
     entity_description: TeslemetrySwitchEntityDescription
 
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn on the Switch."""
+        self.raise_for_scope(self.entity_description.scopes[0])
+        await handle_vehicle_command(self.entity_description.on_func(self.api))
+        self._attr_is_on = True
+        self.async_write_ha_state()
 
-class TeslemetryVehicleSwitchEntity(TeslemetryVehicleEntity, TeslemetrySwitchEntity):
+    async def async_turn_off(self, **kwargs: Any) -> None:
-    """Base class for Teslemetry vehicle switch entities."""
+        """Turn off the Switch."""
+        self.raise_for_scope(self.entity_description.scopes[0])
+        await handle_vehicle_command(self.entity_description.off_func(self.api))
+        self._attr_is_on = False
+        self.async_write_ha_state()
+
+
+class TeslemetryPollingVehicleSwitchEntity(
+    TeslemetryVehicleEntity, TeslemetryVehicleSwitchEntity
+):
+    """Base class for Teslemetry polling vehicle switch entities."""
 
     def __init__(
         self,
@@ -151,30 +193,63 @@ class TeslemetryVehicleSwitchEntity(TeslemetryVehicleEntity, TeslemetrySwitchEnt
 
     def _async_update_attrs(self) -> None:
         """Update the attributes of the sensor."""
-        self._attr_is_on = self.entity_description.value_func(self._value)
+        self._attr_is_on = (
-
+            None
-    async def async_turn_on(self, **kwargs: Any) -> None:
+            if self._value is None
-        """Turn on the Switch."""
+            else self.entity_description.value_func(self._value)
-        self.raise_for_scope(self.entity_description.scopes[0])
+        )
-        await self.wake_up_if_asleep()
-        await handle_vehicle_command(self.entity_description.on_func(self.api))
-        self._attr_is_on = True
-        self.async_write_ha_state()
-
-    async def async_turn_off(self, **kwargs: Any) -> None:
-        """Turn off the Switch."""
-        self.raise_for_scope(self.entity_description.scopes[0])
-        await self.wake_up_if_asleep()
-        await handle_vehicle_command(self.entity_description.off_func(self.api))
-        self._attr_is_on = False
-        self.async_write_ha_state()
 
 
-class TeslemetryChargeFromGridSwitchEntity(
+class TeslemetryStreamingVehicleSwitchEntity(
-    TeslemetryEnergyInfoEntity, TeslemetrySwitchEntity
+    TeslemetryVehicleStreamEntity, TeslemetryVehicleSwitchEntity, RestoreEntity
 ):
+    """Base class for Teslemetry streaming vehicle switch entities."""
+
+    def __init__(
+        self,
+        data: TeslemetryVehicleData,
+        description: TeslemetrySwitchEntityDescription,
+        scopes: list[Scope],
+    ) -> None:
+        """Initialize the Switch."""
+
+        self.entity_description = description
+        self.scoped = any(scope in scopes for scope in description.scopes)
+        super().__init__(data, description.key)
+        if description.unique_id:
+            self._attr_unique_id = f"{data.vin}-{description.unique_id}"
+
+    async def async_added_to_hass(self) -> None:
+        """Handle entity which will be added."""
+        await super().async_added_to_hass()
+
+        # Restore previous state
+        if (state := await self.async_get_last_state()) is not None:
+            if state.state == "on":
+                self._attr_is_on = True
+            elif state.state == "off":
+                self._attr_is_on = False
+
+        # Add listener
+        self.async_on_remove(
+            self.entity_description.streaming_listener(
+                self.vehicle.stream_vehicle, self._value_callback
+            )
+        )
+
+    def _value_callback(self, value: StateType) -> None:
+        """Update the value of the entity."""
+        self._attr_is_on = (
+            None if value is None else self.entity_description.streaming_value_fn(value)
+        )
+        self.async_write_ha_state()
+
+
+class TeslemetryChargeFromGridSwitchEntity(TeslemetryEnergyInfoEntity, SwitchEntity):
     """Entity class for Charge From Grid switch."""
 
+    _attr_device_class = SwitchDeviceClass.SWITCH
+
     def __init__(
         self,
         data: TeslemetryEnergyData,
@@ -215,11 +290,11 @@ class TeslemetryChargeFromGridSwitchEntity(
         self.async_write_ha_state()
 
 
-class TeslemetryStormModeSwitchEntity(
+class TeslemetryStormModeSwitchEntity(TeslemetryEnergyInfoEntity, SwitchEntity):
-    TeslemetryEnergyInfoEntity, TeslemetrySwitchEntity
-):
     """Entity class for Storm Mode switch."""
 
+    _attr_device_class = SwitchDeviceClass.SWITCH
+
     def __init__(
         self,
         data: TeslemetryEnergyData,

tests/components/teslemetry/snapshots/test_switch.ambr

@@ -495,3 +495,21 @@
     'state': 'off',
   })
 # ---
+# name: test_switch_streaming[switch.test_auto_seat_climate_left]
+  'on'
+# ---
+# name: test_switch_streaming[switch.test_auto_seat_climate_right]
+  'off'
+# ---
+# name: test_switch_streaming[switch.test_auto_steering_wheel_heater]
+  'on'
+# ---
+# name: test_switch_streaming[switch.test_charge]
+  'on'
+# ---
+# name: test_switch_streaming[switch.test_defrost]
+  'off'
+# ---
+# name: test_switch_streaming[switch.test_sentry_mode]
+  'on'
+# ---

tests/components/teslemetry/test_switch.py

@@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, patch
 
 import pytest
 from syrupy.assertion import SnapshotAssertion
+from teslemetry_stream import Signal
 
 from homeassistant.components.switch import (
     DOMAIN as SWITCH_DOMAIN,
@@ -14,7 +15,7 @@ from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 
-from . import assert_entities, assert_entities_alt, setup_platform
+from . import assert_entities, assert_entities_alt, reload_platform, setup_platform
 from .const import COMMAND_OK, VEHICLE_DATA_ALT
 
 
@@ -22,6 +23,7 @@ async def test_switch(
     hass: HomeAssistant,
     snapshot: SnapshotAssertion,
     entity_registry: er.EntityRegistry,
+    mock_legacy: AsyncMock,
 ) -> None:
     """Tests that the switch entities are correct."""
 
@@ -34,6 +36,7 @@ async def test_switch_alt(
     snapshot: SnapshotAssertion,
     entity_registry: er.EntityRegistry,
     mock_vehicle_data: AsyncMock,
+    mock_legacy: AsyncMock,
 ) -> None:
     """Tests that the switch entities are correct."""
 
@@ -119,3 +122,47 @@ async def test_switch_services(
         state = hass.states.get(entity_id)
         assert state.state == STATE_OFF
         call.assert_called_once()
+
+
+async def test_switch_streaming(
+    hass: HomeAssistant,
+    snapshot: SnapshotAssertion,
+    entity_registry: er.EntityRegistry,
+    mock_vehicle_data: AsyncMock,
+    mock_add_listener: AsyncMock,
+) -> None:
+    """Tests that the switch entities with streaming are correct."""
+
+    entry = await setup_platform(hass, [Platform.SWITCH])
+
+    # Stream update
+    mock_add_listener.send(
+        {
+            "vin": VEHICLE_DATA_ALT["response"]["vin"],
+            "data": {
+                Signal.SENTRY_MODE: "SentryModeStateIdle",
+                Signal.AUTO_SEAT_CLIMATE_LEFT: True,
+                Signal.AUTO_SEAT_CLIMATE_RIGHT: False,
+                Signal.HVAC_STEERING_WHEEL_HEAT_AUTO: True,
+                Signal.DEFROST_MODE: "DefrostModeStateOff",
+                Signal.DETAILED_CHARGE_STATE: "DetailedChargeStateCharging",
+            },
+            "createdAt": "2024-10-04T10:45:17.537Z",
+        }
+    )
+    await hass.async_block_till_done()
+
+    # Reload the entry
+    await reload_platform(hass, entry, [Platform.SWITCH])
+
+    # Assert the entities restored their values
+    for entity_id in (
+        "switch.test_sentry_mode",
+        "switch.test_auto_seat_climate_left",
+        "switch.test_auto_seat_climate_right",
+        "switch.test_auto_steering_wheel_heater",
+        "switch.test_defrost",
+        "switch.test_charge",
+    ):
+        state = hass.states.get(entity_id)
+        assert state.state == snapshot(name=entity_id)