Add streaming to Teslemetry lock platform (#136037)

homeassistant/components/teslemetry/lock.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from itertools import chain
 from typing import Any
 
 from tesla_fleet_api.const import Scope
@@ -10,10 +11,15 @@ from homeassistant.components.lock import LockEntity
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ServiceValidationError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.restore_state import RestoreEntity
 
 from . import TeslemetryConfigEntry
 from .const import DOMAIN
-from .entity import TeslemetryVehicleEntity
+from .entity import (
+    TeslemetryRootEntity,
+    TeslemetryVehicleEntity,
+    TeslemetryVehicleStreamEntity,
+)
 from .helpers import handle_vehicle_command
 from .models import TeslemetryVehicleData
 
@@ -30,31 +36,38 @@ async def async_setup_entry(
     """Set up the Teslemetry lock platform from a config entry."""
 
     async_add_entities(
-        klass(vehicle, Scope.VEHICLE_CMDS in entry.runtime_data.scopes)
+        chain(
-        for klass in (
+            (
-            TeslemetryVehicleLockEntity,
+                TeslemetryPollingVehicleLockEntity(
-            TeslemetryCableLockEntity,
+                    vehicle, Scope.VEHICLE_CMDS in entry.runtime_data.scopes
+                )
+                if vehicle.api.pre2021 or vehicle.firmware < "2024.26"
+                else TeslemetryStreamingVehicleLockEntity(
+                    vehicle, Scope.VEHICLE_CMDS in entry.runtime_data.scopes
+                )
+                for vehicle in entry.runtime_data.vehicles
+            ),
+            (
+                TeslemetryPollingCableLockEntity(
+                    vehicle, Scope.VEHICLE_CMDS in entry.runtime_data.scopes
+                )
+                if vehicle.api.pre2021 or vehicle.firmware < "2024.26"
+                else TeslemetryStreamingCableLockEntity(
+                    vehicle, Scope.VEHICLE_CMDS in entry.runtime_data.scopes
+                )
+                for vehicle in entry.runtime_data.vehicles
+            ),
         )
-        for vehicle in entry.runtime_data.vehicles
     )
 
 
-class TeslemetryVehicleLockEntity(TeslemetryVehicleEntity, LockEntity):
+class TeslemetryVehicleLockEntity(TeslemetryRootEntity, LockEntity):
-    """Lock entity for Teslemetry."""
+    """Base vehicle lock entity for Teslemetry."""
-
-    def __init__(self, data: TeslemetryVehicleData, scoped: bool) -> None:
-        """Initialize the lock."""
-        super().__init__(data, "vehicle_state_locked")
-        self.scoped = scoped
-
-    def _async_update_attrs(self) -> None:
-        """Update entity attributes."""
-        self._attr_is_locked = self._value
 
     async def async_lock(self, **kwargs: Any) -> None:
         """Lock the doors."""
         self.raise_for_scope(Scope.VEHICLE_CMDS)
-        await self.wake_up_if_asleep()
+
         await handle_vehicle_command(self.api.door_lock())
         self._attr_is_locked = True
         self.async_write_ha_state()
@@ -62,27 +75,65 @@ class TeslemetryVehicleLockEntity(TeslemetryVehicleEntity, LockEntity):
     async def async_unlock(self, **kwargs: Any) -> None:
         """Unlock the doors."""
         self.raise_for_scope(Scope.VEHICLE_CMDS)
-        await self.wake_up_if_asleep()
+
         await handle_vehicle_command(self.api.door_unlock())
         self._attr_is_locked = False
         self.async_write_ha_state()
 
 
-class TeslemetryCableLockEntity(TeslemetryVehicleEntity, LockEntity):
+class TeslemetryPollingVehicleLockEntity(
-    """Cable Lock entity for Teslemetry."""
+    TeslemetryVehicleEntity, TeslemetryVehicleLockEntity
+):
+    """Polling vehicle lock entity for Teslemetry."""
 
-    def __init__(
+    def __init__(self, data: TeslemetryVehicleData, scoped: bool) -> None:
-        self,
+        """Initialize the sensor."""
-        data: TeslemetryVehicleData,
+        super().__init__(
-        scoped: bool,
+            data,
-    ) -> None:
+            "vehicle_state_locked",
-        """Initialize the lock."""
+        )
-        super().__init__(data, "charge_state_charge_port_latch")
         self.scoped = scoped
 
     def _async_update_attrs(self) -> None:
         """Update entity attributes."""
-        self._attr_is_locked = self._value == ENGAGED
+        self._attr_is_locked = self._value
+
+
+class TeslemetryStreamingVehicleLockEntity(
+    TeslemetryVehicleStreamEntity, TeslemetryVehicleLockEntity, RestoreEntity
+):
+    """Streaming vehicle lock entity for Teslemetry."""
+
+    def __init__(self, data: TeslemetryVehicleData, scoped: bool) -> None:
+        """Initialize the sensor."""
+        super().__init__(
+            data,
+            "vehicle_state_locked",
+        )
+        self.scoped = scoped
+
+    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 == "locked":
+                self._attr_is_locked = True
+            elif state.state == "unlocked":
+                self._attr_is_locked = False
+
+        # Add streaming listener
+        self.async_on_remove(self.vehicle.stream_vehicle.listen_Locked(self._callback))
+
+    def _callback(self, value: bool | None) -> None:
+        """Update entity attributes."""
+        self._attr_is_locked = value
+        self.async_write_ha_state()
+
+
+class TeslemetryCableLockEntity(TeslemetryRootEntity, LockEntity):
+    """Base cable Lock entity for Teslemetry."""
 
     async def async_lock(self, **kwargs: Any) -> None:
         """Charge cable Lock cannot be manually locked."""
@@ -95,7 +146,70 @@ class TeslemetryCableLockEntity(TeslemetryVehicleEntity, LockEntity):
     async def async_unlock(self, **kwargs: Any) -> None:
         """Unlock charge cable lock."""
         self.raise_for_scope(Scope.VEHICLE_CMDS)
-        await self.wake_up_if_asleep()
+
         await handle_vehicle_command(self.api.charge_port_door_open())
         self._attr_is_locked = False
         self.async_write_ha_state()
+
+
+class TeslemetryPollingCableLockEntity(
+    TeslemetryVehicleEntity, TeslemetryCableLockEntity
+):
+    """Polling cable lock entity for Teslemetry."""
+
+    def __init__(
+        self,
+        data: TeslemetryVehicleData,
+        scoped: bool,
+    ) -> None:
+        """Initialize the sensor."""
+        super().__init__(
+            data,
+            "charge_state_charge_port_latch",
+        )
+        self.scoped = scoped
+
+    def _async_update_attrs(self) -> None:
+        """Update entity attributes."""
+        if self._value is None:
+            self._attr_is_locked = None
+        self._attr_is_locked = self._value == ENGAGED
+
+
+class TeslemetryStreamingCableLockEntity(
+    TeslemetryVehicleStreamEntity, TeslemetryCableLockEntity, RestoreEntity
+):
+    """Streaming cable lock entity for Teslemetry."""
+
+    def __init__(
+        self,
+        data: TeslemetryVehicleData,
+        scoped: bool,
+    ) -> None:
+        """Initialize the sensor."""
+        super().__init__(
+            data,
+            "charge_state_charge_port_latch",
+        )
+        self.scoped = scoped
+
+    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 == "locked":
+                self._attr_is_locked = True
+            elif state.state == "unlocked":
+                self._attr_is_locked = False
+
+        # Add streaming listener
+        self.async_on_remove(
+            self.vehicle.stream_vehicle.listen_ChargePortLatch(self._callback)
+        )
+
+    def _callback(self, value: str | None) -> None:
+        """Update entity attributes."""
+        self._attr_is_locked = None if value is None else value == ENGAGED
+        self.async_write_ha_state()

tests/components/teslemetry/__init__.py

@@ -14,7 +14,9 @@ from .const import CONFIG
 from tests.common import MockConfigEntry
 
 
-async def setup_platform(hass: HomeAssistant, platforms: list[Platform] | None = None):
+async def setup_platform(
+    hass: HomeAssistant, platforms: list[Platform] | None = None
+) -> MockConfigEntry:
     """Set up the Teslemetry platform."""
 
     mock_entry = MockConfigEntry(

tests/components/teslemetry/fixtures/vehicle_data_alt.json

@@ -35,7 +35,7 @@
       "charge_port_cold_weather_mode": false,
       "charge_port_color": "<invalid>",
       "charge_port_door_open": true,
-      "charge_port_latch": "Engaged",
+      "charge_port_latch": null,
       "charge_rate": 0,
       "charger_actual_current": 0,
       "charger_phases": null,

tests/components/teslemetry/snapshots/test_lock.ambr

@@ -93,3 +93,109 @@
     'state': 'unlocked',
   })
 # ---
+# name: test_lock_alt[lock.test_charge_cable_lock-entry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': None,
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'lock',
+    'entity_category': None,
+    'entity_id': 'lock.test_charge_cable_lock',
+    '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': 'Charge cable lock',
+    'platform': 'teslemetry',
+    'previous_unique_id': None,
+    'supported_features': 0,
+    'translation_key': 'charge_state_charge_port_latch',
+    'unique_id': 'LRW3F7EK4NC700000-charge_state_charge_port_latch',
+    'unit_of_measurement': None,
+  })
+# ---
+# name: test_lock_alt[lock.test_charge_cable_lock-state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Test Charge cable lock',
+      'supported_features': <LockEntityFeature: 0>,
+    }),
+    'context': <ANY>,
+    'entity_id': 'lock.test_charge_cable_lock',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'unlocked',
+  })
+# ---
+# name: test_lock_alt[lock.test_lock-entry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': None,
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'lock',
+    'entity_category': None,
+    'entity_id': 'lock.test_lock',
+    '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': 'Lock',
+    'platform': 'teslemetry',
+    'previous_unique_id': None,
+    'supported_features': 0,
+    'translation_key': 'vehicle_state_locked',
+    'unique_id': 'LRW3F7EK4NC700000-vehicle_state_locked',
+    'unit_of_measurement': None,
+  })
+# ---
+# name: test_lock_alt[lock.test_lock-state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Test Lock',
+      'supported_features': <LockEntityFeature: 0>,
+    }),
+    'context': <ANY>,
+    'entity_id': 'lock.test_lock',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'unlocked',
+  })
+# ---
+# name: test_lock_streaming[lock.test_charge_cable_lock-locked]
+  'locked'
+# ---
+# name: test_lock_streaming[lock.test_charge_cable_lock-unlocked]
+  'unlocked'
+# ---
+# name: test_lock_streaming[lock.test_lock-locked]
+  'locked'
+# ---
+# name: test_lock_streaming[lock.test_lock-unlocked]
+  'unlocked'
+# ---

tests/components/teslemetry/test_lock.py

@@ -1,9 +1,10 @@
 """Test the Teslemetry lock platform."""
 
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
 import pytest
 from syrupy.assertion import SnapshotAssertion
+from teslemetry_stream.const import Signal
 
 from homeassistant.components.lock import (
     DOMAIN as LOCK_DOMAIN,
@@ -16,14 +17,15 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ServiceValidationError
 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
+from .const import COMMAND_OK, VEHICLE_DATA_ALT
 
 
 async def test_lock(
     hass: HomeAssistant,
     snapshot: SnapshotAssertion,
     entity_registry: er.EntityRegistry,
+    mock_legacy: AsyncMock,
 ) -> None:
     """Tests that the lock entities are correct."""
 
@@ -31,6 +33,20 @@ async def test_lock(
     assert_entities(hass, entry.entry_id, entity_registry, snapshot)
 
 
+async def test_lock_alt(
+    hass: HomeAssistant,
+    snapshot: SnapshotAssertion,
+    entity_registry: er.EntityRegistry,
+    mock_vehicle_data: AsyncMock,
+    mock_legacy: AsyncMock,
+) -> None:
+    """Tests that the lock entities are correct."""
+
+    mock_vehicle_data.return_value = VEHICLE_DATA_ALT
+    entry = await setup_platform(hass, [Platform.LOCK])
+    assert_entities(hass, entry.entry_id, entity_registry, snapshot)
+
+
 async def test_lock_services(
     hass: HomeAssistant,
 ) -> None:
@@ -91,3 +107,60 @@ async def test_lock_services(
         state = hass.states.get(entity_id)
         assert state.state == LockState.UNLOCKED
         call.assert_called_once()
+
+
+async def test_lock_streaming(
+    hass: HomeAssistant,
+    snapshot: SnapshotAssertion,
+    mock_vehicle_data: AsyncMock,
+    mock_add_listener: AsyncMock,
+) -> None:
+    """Tests that the lock entities with streaming are correct."""
+
+    entry = await setup_platform(hass, [Platform.LOCK])
+
+    # Stream update
+    mock_add_listener.send(
+        {
+            "vin": VEHICLE_DATA_ALT["response"]["vin"],
+            "data": {
+                Signal.LOCKED: True,
+                Signal.CHARGE_PORT_LATCH: "ChargePortLatchEngaged",
+            },
+            "createdAt": "2024-10-04T10:45:17.537Z",
+        }
+    )
+    await hass.async_block_till_done()
+
+    await reload_platform(hass, entry, [Platform.LOCK])
+
+    # Assert the entities restored their values
+    for entity_id in (
+        "lock.test_lock",
+        "lock.test_charge_cable_lock",
+    ):
+        state = hass.states.get(entity_id)
+        assert state.state == snapshot(name=f"{entity_id}-locked")
+
+    # Stream update
+    mock_add_listener.send(
+        {
+            "vin": VEHICLE_DATA_ALT["response"]["vin"],
+            "data": {
+                Signal.LOCKED: False,
+                Signal.CHARGE_PORT_LATCH: "ChargePortLatchDisengaged",
+            },
+            "createdAt": "2024-10-04T10:45:17.537Z",
+        }
+    )
+    await hass.async_block_till_done()
+
+    await reload_platform(hass, entry, [Platform.LOCK])
+
+    # Assert the entities restored their values
+    for entity_id in (
+        "lock.test_lock",
+        "lock.test_charge_cable_lock",
+    ):
+        state = hass.states.get(entity_id)
+        assert state.state == snapshot(name=f"{entity_id}-unlocked")