Support restoring SensorEntity native_value (#66068)

tests/components/sensor/test_init.py

@@ -11,10 +11,14 @@ from homeassistant.const import (
     TEMP_CELSIUS,
     TEMP_FAHRENHEIT,
 )
+from homeassistant.core import State
+from homeassistant.helpers.restore_state import STORAGE_KEY as RESTORE_STATE_KEY
 from homeassistant.setup import async_setup_component
 from homeassistant.util import dt as dt_util
 from homeassistant.util.unit_system import IMPERIAL_SYSTEM, METRIC_SYSTEM
 
+from tests.common import mock_restore_cache_with_extra_data
+
 
 @pytest.mark.parametrize(
     "unit_system,native_unit,state_unit,native_value,state_value",
@@ -210,3 +214,131 @@ async def test_reject_timezoneless_datetime_str(
         "Invalid datetime: sensor.test provides state '2017-12-19 18:29:42', "
         "which is missing timezone information"
     ) in caplog.text
+
+
+RESTORE_DATA = {
+    "str": {"native_unit_of_measurement": "°F", "native_value": "abc123"},
+    "int": {"native_unit_of_measurement": "°F", "native_value": 123},
+    "float": {"native_unit_of_measurement": "°F", "native_value": 123.0},
+    "date": {
+        "native_unit_of_measurement": "°F",
+        "native_value": {
+            "__type": "<class 'datetime.date'>",
+            "isoformat": date(2020, 2, 8).isoformat(),
+        },
+    },
+    "datetime": {
+        "native_unit_of_measurement": "°F",
+        "native_value": {
+            "__type": "<class 'datetime.datetime'>",
+            "isoformat": datetime(2020, 2, 8, 15, tzinfo=timezone.utc).isoformat(),
+        },
+    },
+}
+
+
+# None | str | int | float | date | datetime:
+@pytest.mark.parametrize(
+    "native_value, native_value_type, expected_extra_data, device_class",
+    [
+        ("abc123", str, RESTORE_DATA["str"], None),
+        (123, int, RESTORE_DATA["int"], SensorDeviceClass.TEMPERATURE),
+        (123.0, float, RESTORE_DATA["float"], SensorDeviceClass.TEMPERATURE),
+        (date(2020, 2, 8), dict, RESTORE_DATA["date"], SensorDeviceClass.DATE),
+        (
+            datetime(2020, 2, 8, 15, tzinfo=timezone.utc),
+            dict,
+            RESTORE_DATA["datetime"],
+            SensorDeviceClass.TIMESTAMP,
+        ),
+    ],
+)
+async def test_restore_sensor_save_state(
+    hass,
+    enable_custom_integrations,
+    hass_storage,
+    native_value,
+    native_value_type,
+    expected_extra_data,
+    device_class,
+):
+    """Test RestoreSensor."""
+    platform = getattr(hass.components, "test.sensor")
+    platform.init(empty=True)
+    platform.ENTITIES["0"] = platform.MockRestoreSensor(
+        name="Test",
+        native_value=native_value,
+        native_unit_of_measurement=TEMP_FAHRENHEIT,
+        device_class=device_class,
+    )
+
+    entity0 = platform.ENTITIES["0"]
+    assert await async_setup_component(hass, "sensor", {"sensor": {"platform": "test"}})
+    await hass.async_block_till_done()
+
+    # Trigger saving state
+    await hass.async_stop()
+
+    assert len(hass_storage[RESTORE_STATE_KEY]["data"]) == 1
+    state = hass_storage[RESTORE_STATE_KEY]["data"][0]["state"]
+    assert state["entity_id"] == entity0.entity_id
+    extra_data = hass_storage[RESTORE_STATE_KEY]["data"][0]["extra_data"]
+    assert extra_data == expected_extra_data
+    assert type(extra_data["native_value"]) == native_value_type
+
+
+@pytest.mark.parametrize(
+    "native_value, native_value_type, extra_data, device_class, uom",
+    [
+        ("abc123", str, RESTORE_DATA["str"], None, "°F"),
+        (123, int, RESTORE_DATA["int"], SensorDeviceClass.TEMPERATURE, "°F"),
+        (123.0, float, RESTORE_DATA["float"], SensorDeviceClass.TEMPERATURE, "°F"),
+        (date(2020, 2, 8), date, RESTORE_DATA["date"], SensorDeviceClass.DATE, "°F"),
+        (
+            datetime(2020, 2, 8, 15, tzinfo=timezone.utc),
+            datetime,
+            RESTORE_DATA["datetime"],
+            SensorDeviceClass.TIMESTAMP,
+            "°F",
+        ),
+        (None, type(None), None, None, None),
+        (None, type(None), {}, None, None),
+        (None, type(None), {"beer": 123}, None, None),
+        (
+            None,
+            type(None),
+            {"native_unit_of_measurement": "°F", "native_value": {}},
+            None,
+            None,
+        ),
+    ],
+)
+async def test_restore_sensor_restore_state(
+    hass,
+    enable_custom_integrations,
+    hass_storage,
+    native_value,
+    native_value_type,
+    extra_data,
+    device_class,
+    uom,
+):
+    """Test RestoreSensor."""
+    mock_restore_cache_with_extra_data(hass, ((State("sensor.test", ""), extra_data),))
+
+    platform = getattr(hass.components, "test.sensor")
+    platform.init(empty=True)
+    platform.ENTITIES["0"] = platform.MockRestoreSensor(
+        name="Test",
+        device_class=device_class,
+    )
+
+    entity0 = platform.ENTITIES["0"]
+    assert await async_setup_component(hass, "sensor", {"sensor": {"platform": "test"}})
+    await hass.async_block_till_done()
+
+    assert hass.states.get(entity0.entity_id)
+
+    assert entity0.native_value == native_value
+    assert type(entity0.native_value) == native_value_type
+    assert entity0.native_unit_of_measurement == uom