Clean up SensorRestore in Tesla Fleet (#129116)

* Remove, fix, and test restore

* slightly better comment

* use restore instead

* parametrize test

* Apply suggestions from code review

* revert change to Teslemetry

* revert change to Teslemetry

---------

Co-authored-by: G Johansson <goran.johansson@shiftit.se>

homeassistant/components/tesla_fleet/sensor.py

@@ -486,7 +486,7 @@ class TeslaFleetVehicleSensorEntity(TeslaFleetVehicleEntity, RestoreSensor):
     async def async_added_to_hass(self) -> None:
         """Handle entity which will be added."""
         await super().async_added_to_hass()
-        if self.coordinator.data.get("state") == TeslaFleetState.OFFLINE:
+        if self.coordinator.data.get("state") != TeslaFleetState.ONLINE:
             if (sensor_data := await self.async_get_last_sensor_data()) is not None:
                 self._attr_native_value = sensor_data.native_value
 
@@ -524,7 +524,7 @@ class TeslaFleetVehicleTimeSensorEntity(TeslaFleetVehicleEntity, SensorEntity):
             self._attr_native_value = self._get_timestamp(self._value)
 
 
-class TeslaFleetEnergyLiveSensorEntity(TeslaFleetEnergyLiveEntity, RestoreSensor):
+class TeslaFleetEnergyLiveSensorEntity(TeslaFleetEnergyLiveEntity, SensorEntity):
     """Base class for Tesla Fleet energy site metric sensors."""
 
     entity_description: SensorEntityDescription
@@ -538,20 +538,13 @@ class TeslaFleetEnergyLiveSensorEntity(TeslaFleetEnergyLiveEntity, RestoreSensor
         self.entity_description = description
         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()
-        if not self.coordinator.updated_once:
-            if (sensor_data := await self.async_get_last_sensor_data()) is not None:
-                self._attr_native_value = sensor_data.native_value
-
     def _async_update_attrs(self) -> None:
         """Update the attributes of the sensor."""
         self._attr_available = not self.is_none
         self._attr_native_value = self._value
 
 
-class TeslaFleetWallConnectorSensorEntity(TeslaFleetWallConnectorEntity, RestoreSensor):
+class TeslaFleetWallConnectorSensorEntity(TeslaFleetWallConnectorEntity, SensorEntity):
     """Base class for Tesla Fleet energy site metric sensors."""
 
     entity_description: SensorEntityDescription
@@ -570,20 +563,13 @@ class TeslaFleetWallConnectorSensorEntity(TeslaFleetWallConnectorEntity, Restore
             description.key,
         )
 
-    async def async_added_to_hass(self) -> None:
-        """Handle entity which will be added."""
-        await super().async_added_to_hass()
-        if not self.coordinator.updated_once:
-            if (sensor_data := await self.async_get_last_sensor_data()) is not None:
-                self._attr_native_value = sensor_data.native_value
-
     def _async_update_attrs(self) -> None:
         """Update the attributes of the sensor."""
         self._attr_available = not self.is_none
         self._attr_native_value = self._value
 
 
-class TeslaFleetEnergyInfoSensorEntity(TeslaFleetEnergyInfoEntity, RestoreSensor):
+class TeslaFleetEnergyInfoSensorEntity(TeslaFleetEnergyInfoEntity, SensorEntity):
     """Base class for Tesla Fleet energy site metric sensors."""
 
     entity_description: SensorEntityDescription
@@ -597,13 +583,6 @@ class TeslaFleetEnergyInfoSensorEntity(TeslaFleetEnergyInfoEntity, RestoreSensor
         self.entity_description = description
         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()
-        if not self.coordinator.updated_once:
-            if (sensor_data := await self.async_get_last_sensor_data()) is not None:
-                self._attr_native_value = sensor_data.native_value
-
     def _async_update_attrs(self) -> None:
         """Update the attributes of the sensor."""
         self._attr_available = not self.is_none

tests/components/tesla_fleet/test_sensor.py

@@ -1,13 +1,14 @@
 """Test the Tesla Fleet sensor platform."""
 
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, patch
 
 from freezegun.api import FrozenDateTimeFactory
 import pytest
 from syrupy.assertion import SnapshotAssertion
+from tesla_fleet_api.exceptions import VehicleOffline
 
 from homeassistant.components.tesla_fleet.coordinator import VEHICLE_INTERVAL
-from homeassistant.const import Platform
+from homeassistant.const import STATE_UNAVAILABLE, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 
@@ -41,3 +42,38 @@ async def test_sensors(
     await hass.async_block_till_done()
 
     assert_entities_alt(hass, normal_config_entry.entry_id, entity_registry, snapshot)
+
+
+@pytest.mark.parametrize(
+    ("entity_id", "initial", "restored"),
+    [
+        ("sensor.test_battery_level", "77", "77"),
+        ("sensor.test_outside_temperature", "30", "30"),
+        ("sensor.test_time_to_arrival", "2024-01-01T00:00:06+00:00", STATE_UNAVAILABLE),
+    ],
+)
+async def test_sensors_restore(
+    hass: HomeAssistant,
+    snapshot: SnapshotAssertion,
+    normal_config_entry: MockConfigEntry,
+    entity_registry: er.EntityRegistry,
+    freezer: FrozenDateTimeFactory,
+    mock_vehicle_data: AsyncMock,
+    entity_id: str,
+    initial: str,
+    restored: str,
+) -> None:
+    """Test if the sensor should restore it's state or not when vehicle is offline."""
+
+    freezer.move_to("2024-01-01 00:00:00+00:00")
+
+    await setup_platform(hass, normal_config_entry, [Platform.SENSOR])
+
+    assert hass.states.get(entity_id).state == initial
+
+    mock_vehicle_data.side_effect = VehicleOffline
+
+    with patch("homeassistant.components.tesla_fleet.PLATFORMS", [Platform.SENSOR]):
+        assert await hass.config_entries.async_reload(normal_config_entry.entry_id)
+
+    assert hass.states.get(entity_id).state == restored