Improve saved state of RestoreSensor when using freezegun

tests/components/conftest.py

@@ -23,7 +23,7 @@ import pytest
 import voluptuous as vol
 
 from homeassistant import components, loader
-from homeassistant.components import repairs
+from homeassistant.components import repairs, sensor
 from homeassistant.config_entries import (
     DISCOVERY_SOURCES,
     ConfigEntriesFlowManager,
@@ -59,6 +59,30 @@ if TYPE_CHECKING:
 RE_REQUEST_DOMAIN = re.compile(r".*tests\/components\/([^/]+)\/.*")
 
 
+@pytest.fixture(scope="session", autouse=True)
+def patch_restore_sensor() -> Generator[None]:
+    """Fixture to patch restore sensor."""
+    real_as_dict = sensor.SensorExtraStoredData.as_dict
+
+    def as_dict(self) -> dict[str, Any]:
+        """Return a dict representation of the sensor data."""
+        data = real_as_dict(self)
+        if not (_type := data["native_value"].get("__type")):
+            return data
+        if _type == "<class 'freezegun.api.FakeDate'>":
+            data["native_value"]["__type"] = "<class 'datetime.date'>"
+        elif _type == "<class 'tests.patch_time.HAFakeDatetime'>":
+            data["native_value"]["__type"] = "<class 'datetime.datetime'>"
+        return data
+
+    with patch(
+        "homeassistant.components.sensor.SensorExtraStoredData.as_dict",
+        autospec=True,
+        side_effect=as_dict,
+    ):
+        yield
+
+
 @pytest.fixture(scope="session", autouse=find_spec("zeroconf") is not None)
 def patch_zeroconf_multiple_catcher() -> Generator[None]:
     """If installed, patch zeroconf wrapper that detects if multiple instances are used."""

tests/components/sensor/test_init.py

@@ -8,6 +8,7 @@ from decimal import Decimal
 from typing import Any
 from unittest.mock import patch
 
+from freezegun.api import freeze_time
 import pytest
 
 from homeassistant.components import sensor
@@ -475,6 +476,62 @@ async def test_restore_sensor_save_state(
     assert type(extra_data["native_value"]) is native_value_type
 
 
+@freeze_time("2020-02-08 15:00:00")
+async def test_restore_sensor_save_state_frozen_time_datetime(
+    hass: HomeAssistant,
+    hass_storage: dict[str, Any],
+) -> None:
+    """Test RestoreSensor."""
+    entity0 = MockRestoreSensor(
+        name="Test",
+        native_value=dt_util.utcnow(),
+        native_unit_of_measurement=None,
+        device_class=SensorDeviceClass.TIMESTAMP,
+    )
+    setup_test_component_platform(hass, sensor.DOMAIN, [entity0])
+
+    assert await async_setup_component(hass, "sensor", {"sensor": {"platform": "test"}})
+    await hass.async_block_till_done()
+
+    # Trigger saving state
+    await async_mock_restore_state_shutdown_restart(hass)
+
+    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 == RESTORE_DATA["datetime"]
+    assert type(extra_data["native_value"]) is dict
+
+
+@freeze_time("2020-02-08 15:00:00")
+async def test_restore_sensor_save_state_frozen_time_date(
+    hass: HomeAssistant,
+    hass_storage: dict[str, Any],
+) -> None:
+    """Test RestoreSensor."""
+    entity0 = MockRestoreSensor(
+        name="Test",
+        native_value=dt_util.utcnow().date(),
+        native_unit_of_measurement=None,
+        device_class=SensorDeviceClass.DATE,
+    )
+    setup_test_component_platform(hass, sensor.DOMAIN, [entity0])
+
+    assert await async_setup_component(hass, "sensor", {"sensor": {"platform": "test"}})
+    await hass.async_block_till_done()
+
+    # Trigger saving state
+    await async_mock_restore_state_shutdown_restart(hass)
+
+    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 == RESTORE_DATA["date"]
+    assert type(extra_data["native_value"]) is dict
+
+
 @pytest.mark.parametrize(
     ("native_value", "native_value_type", "extra_data", "device_class", "uom"),
     [