From dffdc78915ad9d25f54be90ef62659b2c68de347 Mon Sep 17 00:00:00 2001 From: Kevin Stillhammer Date: Wed, 30 Nov 2022 16:13:11 +0100 Subject: [PATCH] Make HERETravelTimeSensor extend RestoreSensor (#82400) --- .../components/here_travel_time/sensor.py | 21 ++- .../here_travel_time/test_sensor.py | 145 +++++++++++++++++- 2 files changed, 157 insertions(+), 9 deletions(-) diff --git a/homeassistant/components/here_travel_time/sensor.py b/homeassistant/components/here_travel_time/sensor.py index e69d1efe4a4..18433a86560 100644 --- a/homeassistant/components/here_travel_time/sensor.py +++ b/homeassistant/components/here_travel_time/sensor.py @@ -6,8 +6,8 @@ from datetime import timedelta from typing import Any from homeassistant.components.sensor import ( + RestoreSensor, SensorDeviceClass, - SensorEntity, SensorEntityDescription, SensorStateClass, ) @@ -99,7 +99,7 @@ async def async_setup_entry( async_add_entities(sensors) -class HERETravelTimeSensor(SensorEntity, CoordinatorEntity): +class HERETravelTimeSensor(RestoreSensor, CoordinatorEntity): """Representation of a HERE travel time sensor.""" def __init__( @@ -121,8 +121,14 @@ class HERETravelTimeSensor(SensorEntity, CoordinatorEntity): ) self._attr_has_entity_name = True + async def _async_restore_state(self) -> None: + """Restore state.""" + if restored_data := await self.async_get_last_sensor_data(): + self._attr_native_value = restored_data.native_value + async def async_added_to_hass(self) -> None: """Wait for start so origin and destination entities can be resolved.""" + await self._async_restore_state() await super().async_added_to_hass() async def _update_at_start(_): @@ -130,12 +136,13 @@ class HERETravelTimeSensor(SensorEntity, CoordinatorEntity): self.async_on_remove(async_at_start(self.hass, _update_at_start)) - @property - def native_value(self) -> str | float | None: - """Return the state of the sensor.""" + def _handle_coordinator_update(self) -> None: + """Handle updated data from the coordinator.""" if self.coordinator.data is not None: - return self.coordinator.data.get(self.entity_description.key) - return None + self._attr_native_value = self.coordinator.data.get( + self.entity_description.key + ) + self.async_write_ha_state() @property def attribution(self) -> str | None: diff --git a/tests/components/here_travel_time/test_sensor.py b/tests/components/here_travel_time/test_sensor.py index 435d24d7b1d..0ad86a04992 100644 --- a/tests/components/here_travel_time/test_sensor.py +++ b/tests/components/here_travel_time/test_sensor.py @@ -34,29 +34,37 @@ from homeassistant.components.here_travel_time.const import ( TRAVEL_MODE_PUBLIC, TRAVEL_MODE_TRUCK, ) +from homeassistant.components.sensor import ( + ATTR_LAST_RESET, + ATTR_STATE_CLASS, + SensorStateClass, +) from homeassistant.const import ( ATTR_ATTRIBUTION, ATTR_ICON, ATTR_LATITUDE, ATTR_LONGITUDE, + ATTR_UNIT_OF_MEASUREMENT, CONF_API_KEY, CONF_MODE, CONF_NAME, EVENT_HOMEASSISTANT_START, TIME_MINUTES, + UnitOfLength, ) -from homeassistant.core import HomeAssistant +from homeassistant.core import CoreState, HomeAssistant, State from homeassistant.setup import async_setup_component from .const import ( API_KEY, + DEFAULT_CONFIG, DESTINATION_LATITUDE, DESTINATION_LONGITUDE, ORIGIN_LATITUDE, ORIGIN_LONGITUDE, ) -from tests.common import MockConfigEntry +from tests.common import MockConfigEntry, mock_restore_cache_with_extra_data @pytest.mark.parametrize( @@ -440,3 +448,136 @@ async def test_route_not_found(hass: HomeAssistant, caplog): await hass.async_block_till_done() assert "Route calculation failed: Couldn't find a route." in caplog.text + + +@pytest.mark.usefixtures("valid_response") +async def test_restore_state(hass): + """Test sensor restore state.""" + # Home assistant is not running yet + hass.state = CoreState.not_running + last_reset = "2022-11-29T00:00:00.000000+00:00" + mock_restore_cache_with_extra_data( + hass, + [ + ( + State( + "sensor.test_duration", + "1234", + attributes={ + ATTR_LAST_RESET: last_reset, + ATTR_UNIT_OF_MEASUREMENT: TIME_MINUTES, + ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT, + }, + ), + { + "native_value": 1234, + "native_unit_of_measurement": TIME_MINUTES, + "icon": "mdi:car", + "last_reset": last_reset, + }, + ), + ( + State( + "sensor.test_duration_in_traffic", + "5678", + attributes={ + ATTR_LAST_RESET: last_reset, + ATTR_UNIT_OF_MEASUREMENT: TIME_MINUTES, + ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT, + }, + ), + { + "native_value": 5678, + "native_unit_of_measurement": TIME_MINUTES, + "icon": "mdi:car", + "last_reset": last_reset, + }, + ), + ( + State( + "sensor.test_distance", + "123", + attributes={ + ATTR_LAST_RESET: last_reset, + ATTR_UNIT_OF_MEASUREMENT: UnitOfLength.KILOMETERS, + ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT, + }, + ), + { + "native_value": 123, + "native_unit_of_measurement": UnitOfLength.KILOMETERS, + "icon": "mdi:car", + "last_reset": last_reset, + }, + ), + ( + State( + "sensor.test_origin", + "Origin Address 1", + attributes={ + ATTR_LAST_RESET: last_reset, + ATTR_LATITUDE: ORIGIN_LATITUDE, + ATTR_LONGITUDE: ORIGIN_LONGITUDE, + }, + ), + { + "native_value": "Origin Address 1", + "native_unit_of_measurement": None, + ATTR_LATITUDE: ORIGIN_LATITUDE, + ATTR_LONGITUDE: ORIGIN_LONGITUDE, + "icon": "mdi:store-marker", + "last_reset": last_reset, + }, + ), + ( + State( + "sensor.test_destination", + "Destination Address 1", + attributes={ + ATTR_LAST_RESET: last_reset, + ATTR_LATITUDE: DESTINATION_LATITUDE, + ATTR_LONGITUDE: DESTINATION_LONGITUDE, + }, + ), + { + "native_value": "Destination Address 1", + "native_unit_of_measurement": None, + "icon": "mdi:store-marker", + "last_reset": last_reset, + }, + ), + ], + ) + + # create and add entry + mock_entry = MockConfigEntry( + domain=DOMAIN, unique_id=DOMAIN, data=DEFAULT_CONFIG, options=DEFAULT_OPTIONS + ) + mock_entry.add_to_hass(hass) + + await hass.config_entries.async_setup(mock_entry.entry_id) + await hass.async_block_till_done() + + print(hass.states.async_all()) + + # restore from cache + state = hass.states.get("sensor.test_duration") + assert state.state == "1234" + assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == TIME_MINUTES + assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT + + state = hass.states.get("sensor.test_duration_in_traffic") + assert state.state == "5678" + assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == TIME_MINUTES + assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT + + state = hass.states.get("sensor.test_distance") + assert state.state == "123" + assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == UnitOfLength.KILOMETERS + assert state.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT + + state = hass.states.get("sensor.test_origin") + assert state.state == "Origin Address 1" + + state = hass.states.get("sensor.test_destination") + assert state.state == "Destination Address 1"