Make HERETravelTimeSensor extend RestoreSensor (#82400)

pull/83436/head
Kevin Stillhammer 2022-11-30 16:13:11 +01:00 committed by GitHub
parent a12568440d
commit dffdc78915
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 157 additions and 9 deletions

View File

@ -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:

View File

@ -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"