Add streaming to Teslemetry number platform (home-assistant#136048)

homeassistant/components/teslemetry/number.py

@@ -9,20 +9,33 @@
 
 from tesla_fleet_api import EnergySpecific, VehicleSpecific
 from tesla_fleet_api.const import Scope
+from teslemetry_stream import TeslemetryStreamVehicle
 
 from homeassistant.components.number import (
     NumberDeviceClass,
     NumberEntity,
     NumberEntityDescription,
     NumberMode,
+    RestoreNumber,
+)
+from homeassistant.const import (
+    PERCENTAGE,
+    PRECISION_WHOLE,
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+    UnitOfElectricCurrent,
 )
-from homeassistant.const import PERCENTAGE, PRECISION_WHOLE, UnitOfElectricCurrent
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.icon import icon_for_battery_level
 
 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
 
@@ -33,12 +46,22 @@
 class TeslemetryNumberVehicleEntityDescription(NumberEntityDescription):
     """Describes Teslemetry Number entity."""
 
-    func: Callable[[VehicleSpecific, float], Awaitable[Any]]
-    native_min_value: float
-    native_max_value: float
+    func: Callable[[VehicleSpecific, int], Awaitable[Any]]
     min_key: str | None = None
     max_key: str
+    native_min_value: float
+    native_max_value: float
     scopes: list[Scope]
+    value_listener: Callable[
+        [TeslemetryStreamVehicle, Callable[[int | None], None]],
+        Callable[[], None],
+    ]
+    max_listener: (
+        Callable[
+            [TeslemetryStreamVehicle, Callable[[int | None], None]], Callable[[], None]
+        ]
+        | None
+    ) = None
 
 
 VEHICLE_DESCRIPTIONS: tuple[TeslemetryNumberVehicleEntityDescription, ...] = (
@@ -52,7 +75,9 @@ class TeslemetryNumberVehicleEntityDescription(NumberEntityDescription):
         mode=NumberMode.AUTO,
         max_key="charge_state_charge_current_request_max",
         func=lambda api, value: api.set_charging_amps(value),
-        scopes=[Scope.VEHICLE_CHARGING_CMDS],
+        scopes=[Scope.VEHICLE_CHARGING_CMDS, Scope.VEHICLE_CMDS],
+        value_listener=lambda x, y: x.listen_ChargeCurrentRequest(y),
+        max_listener=lambda x, y: x.listen_ChargeCurrentRequestMax(y),
     ),
     TeslemetryNumberVehicleEntityDescription(
         key="charge_state_charge_limit_soc",
@@ -62,10 +87,10 @@ class TeslemetryNumberVehicleEntityDescription(NumberEntityDescription):
         native_unit_of_measurement=PERCENTAGE,
         device_class=NumberDeviceClass.BATTERY,
         mode=NumberMode.AUTO,
-        min_key="charge_state_charge_limit_soc_min",
         max_key="charge_state_charge_limit_soc_max",
         func=lambda api, value: api.set_charge_limit(value),
         scopes=[Scope.VEHICLE_CHARGING_CMDS, Scope.VEHICLE_CMDS],
+        value_listener=lambda x, y: x.listen_ChargeLimitSoc(y),
     ),
 )
 
@@ -76,16 +101,29 @@ class TeslemetryNumberBatteryEntityDescription(NumberEntityDescription):
 
     func: Callable[[EnergySpecific, float], Awaitable[Any]]
     requires: str | None = None
+    scopes: list[Scope]
 
 
 ENERGY_INFO_DESCRIPTIONS: tuple[TeslemetryNumberBatteryEntityDescription, ...] = (
     TeslemetryNumberBatteryEntityDescription(
         key="backup_reserve_percent",
+        native_step=PRECISION_WHOLE,
+        native_min_value=0,
+        native_max_value=100,
+        device_class=NumberDeviceClass.BATTERY,
+        native_unit_of_measurement=PERCENTAGE,
+        scopes=[Scope.ENERGY_CMDS],
         func=lambda api, value: api.backup(int(value)),
         requires="components_battery",
     ),
     TeslemetryNumberBatteryEntityDescription(
         key="off_grid_vehicle_charging_reserve_percent",
+        native_step=PRECISION_WHOLE,
+        native_min_value=0,
+        native_max_value=100,
+        device_class=NumberDeviceClass.BATTERY,
+        native_unit_of_measurement=PERCENTAGE,
+        scopes=[Scope.ENERGY_CMDS],
         func=lambda api, value: api.off_grid_vehicle_charging_reserve(int(value)),
         requires="components_off_grid_vehicle_charging_reserve_supported",
     ),
@@ -101,16 +139,22 @@ async def async_setup_entry(
 
     async_add_entities(
         chain(
-            (  # Add vehicle entities
-                TeslemetryVehicleNumberEntity(
+            (
+                TeslemetryPollingNumberEntity(
+                    vehicle,
+                    description,
+                    entry.runtime_data.scopes,
+                )
+                if vehicle.api.pre2021 or vehicle.firmware < "2024.26"
+                else TeslemetryStreamingNumberEntity(
                     vehicle,
                     description,
                     entry.runtime_data.scopes,
                 )
                 for vehicle in entry.runtime_data.vehicles
                 for description in VEHICLE_DESCRIPTIONS
             ),
-            (  # Add energy site entities
+            (
                 TeslemetryEnergyInfoNumberSensorEntity(
                     energysite,
                     description,
@@ -125,11 +169,25 @@ async def async_setup_entry(
     )
 
 
-class TeslemetryVehicleNumberEntity(TeslemetryVehicleEntity, NumberEntity):
+class TeslemetryVehicleNumberEntity(TeslemetryRootEntity, NumberEntity):
     """Vehicle number entity base class."""
 
     entity_description: TeslemetryNumberVehicleEntityDescription
 
+    async def async_set_native_value(self, value: float) -> None:
+        """Set new value."""
+        value = int(value)
+        self.raise_for_scope(self.entity_description.scopes[0])
+        await handle_vehicle_command(self.entity_description.func(self.api, value))
+        self._attr_native_value = value
+        self.async_write_ha_state()
+
+
+class TeslemetryPollingNumberEntity(
+    TeslemetryVehicleEntity, TeslemetryVehicleNumberEntity
+):
+    """Vehicle polling number entity."""
+
     def __init__(
         self,
         data: TeslemetryVehicleData,
@@ -148,26 +206,67 @@ def _async_update_attrs(self) -> None:
         """Update the attributes of the entity."""
         self._attr_native_value = self._value
 
-        if (min_key := self.entity_description.min_key) is not None:
-            self._attr_native_min_value = self.get_number(
-                min_key,
-                self.entity_description.native_min_value,
-            )
-        else:
-            self._attr_native_min_value = self.entity_description.native_min_value
-
         self._attr_native_max_value = self.get_number(
             self.entity_description.max_key,
             self.entity_description.native_max_value,
         )
 
-    async def async_set_native_value(self, value: float) -> None:
-        """Set new value."""
-        value = int(value)
-        self.raise_for_scope(self.entity_description.scopes[0])
-        await self.wake_up_if_asleep()
-        await handle_vehicle_command(self.entity_description.func(self.api, value))
-        self._attr_native_value = value
+
+class TeslemetryStreamingNumberEntity(
+    TeslemetryVehicleStreamEntity, TeslemetryVehicleNumberEntity, RestoreNumber
+):
+    """Number entity for current charge."""
+
+    entity_description: TeslemetryNumberVehicleEntityDescription
+
+    def __init__(
+        self,
+        data: TeslemetryVehicleData,
+        description: TeslemetryNumberVehicleEntityDescription,
+        scopes: list[Scope],
+    ) -> None:
+        """Initialize the Number entity."""
+        self.scoped = any(scope in scopes for scope in description.scopes)
+        self.entity_description = description
+        self._attr_native_max_value = self.entity_description.native_max_value
+        super().__init__(data, description.key)
+
+    async def async_added_to_hass(self) -> None:
+        """Handle entity which will be added."""
+        await super().async_added_to_hass()
+
+        # Restore state
+        if (last_state := await self.async_get_last_state()) and (
+            last_number_data := await self.async_get_last_number_data()
+        ):
+            if last_state.state not in (STATE_UNKNOWN, STATE_UNAVAILABLE):
+                self._attr_native_value = last_number_data.native_value
+            if last_number_data.native_max_value:
+                self._attr_native_max_value = last_number_data.native_max_value
+
+        # Add listeners
+        self.async_on_remove(
+            self.entity_description.value_listener(
+                self.vehicle.stream_vehicle, self._value_callback
+            )
+        )
+        if self.entity_description.max_listener:
+            self.async_on_remove(
+                self.entity_description.max_listener(
+                    self.vehicle.stream_vehicle, self._max_callback
+                )
+            )
+
+    def _value_callback(self, value: int | None) -> None:
+        """Update the value of the entity."""
+        self._attr_native_value = None if value is None else value
+        self.async_write_ha_state()
+
+    def _max_callback(self, value: int | None) -> None:
+        """Update the value of the entity."""
+        self._attr_native_max_value = (
+            self.entity_description.native_max_value if value is None else value
+        )
         self.async_write_ha_state()
 
 

tests/components/teslemetry/__init__.py

@@ -32,6 +32,19 @@ async def setup_platform(hass: HomeAssistant, platforms: list[Platform] | None =
     return mock_entry
 
 
+async def reload_platform(
+    hass: HomeAssistant, entry: MockConfigEntry, platforms: list[Platform] | None = None
+):
+    """Reload the Teslemetry platform."""
+
+    if platforms is None:
+        await hass.config_entries.async_reload(entry.entry_id)
+    else:
+        with patch("homeassistant.components.teslemetry.PLATFORMS", platforms):
+            await hass.config_entries.async_reload(entry.entry_id)
+    await hass.async_block_till_done()
+
+
 def assert_entities(
     hass: HomeAssistant,
     entry_id: str,

tests/components/teslemetry/snapshots/test_number.ambr

@@ -229,3 +229,9 @@
     'state': '80',
   })
 # ---
+# name: test_number_streaming[number.test_charge_current-state]
+  '24'
+# ---
+# name: test_number_streaming[number.test_charge_limit-state]
+  '99'
+# ---

tests/components/teslemetry/test_number.py

@@ -4,6 +4,7 @@
 
 import pytest
 from syrupy.assertion import SnapshotAssertion
+from teslemetry_stream import Signal
 
 from homeassistant.components.number import (
     ATTR_VALUE,
@@ -14,7 +15,7 @@
 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
 
 
@@ -23,6 +24,7 @@ async def test_number(
     hass: HomeAssistant,
     snapshot: SnapshotAssertion,
     entity_registry: er.EntityRegistry,
+    mock_legacy: AsyncMock,
 ) -> None:
     """Tests that the number entities are correct."""
 
@@ -100,3 +102,38 @@ async def test_number_services(
         state = hass.states.get(entity_id)
         assert state.state == "88"
         call.assert_called_once()
+
+
+async def test_number_streaming(
+    hass: HomeAssistant,
+    snapshot: SnapshotAssertion,
+    mock_vehicle_data: AsyncMock,
+    mock_add_listener: AsyncMock,
+) -> None:
+    """Tests that the number entities with streaming are correct."""
+
+    entry = await setup_platform(hass, [Platform.NUMBER])
+
+    # Stream update
+    mock_add_listener.send(
+        {
+            "vin": VEHICLE_DATA_ALT["response"]["vin"],
+            "data": {
+                Signal.CHARGE_CURRENT_REQUEST: 24,
+                Signal.CHARGE_CURRENT_REQUEST_MAX: 32,
+                Signal.CHARGE_LIMIT_SOC: 99,
+            },
+            "createdAt": "2024-10-04T10:45:17.537Z",
+        }
+    )
+    await hass.async_block_till_done()
+
+    await reload_platform(hass, entry, [Platform.NUMBER])
+
+    # Assert the entities restored their values
+    for entity_id in (
+        "number.test_charge_current",
+        "number.test_charge_limit",
+    ):
+        state = hass.states.get(entity_id)
+        assert state.state == snapshot(name=f"{entity_id}-state")