Add Save Persistent States service (#53881)

homeassistant/components/homeassistant/__init__.py

@@ -14,13 +14,14 @@ from homeassistant.const import (
     RESTART_EXIT_CODE,
     SERVICE_HOMEASSISTANT_RESTART,
     SERVICE_HOMEASSISTANT_STOP,
+    SERVICE_SAVE_PERSISTENT_STATES,
     SERVICE_TOGGLE,
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
 )
 import homeassistant.core as ha
 from homeassistant.exceptions import HomeAssistantError, Unauthorized, UnknownUser
-from homeassistant.helpers import config_validation as cv, recorder
+from homeassistant.helpers import config_validation as cv, recorder, restore_state
 from homeassistant.helpers.service import (
     async_extract_config_entry_ids,
     async_extract_referenced_entity_ids,
@@ -53,6 +54,10 @@ SHUTDOWN_SERVICES = (SERVICE_HOMEASSISTANT_STOP, SERVICE_HOMEASSISTANT_RESTART)
 async def async_setup(hass: ha.HomeAssistant, config: dict) -> bool:  # noqa: C901
     """Set up general services related to Home Assistant."""
 
+    async def async_save_persistent_states(service):
+        """Handle calls to homeassistant.save_persistent_states."""
+        await restore_state.RestoreStateData.async_save_persistent_states(hass)
+
     async def async_handle_turn_service(service):
         """Handle calls to homeassistant.turn_on/off."""
         referenced = await async_extract_referenced_entity_ids(hass, service)
@@ -114,6 +119,10 @@ async def async_setup(hass: ha.HomeAssistant, config: dict) -> bool:  # noqa: C9
         if tasks:
             await asyncio.gather(*tasks)
 
+    hass.services.async_register(
+        ha.DOMAIN, SERVICE_SAVE_PERSISTENT_STATES, async_save_persistent_states
+    )
+
     service_schema = vol.Schema({ATTR_ENTITY_ID: cv.entity_ids}, extra=vol.ALLOW_EXTRA)
 
     hass.services.async_register(

homeassistant/components/homeassistant/services.yaml

@@ -74,3 +74,9 @@ reload_config_entry:
       example: 8955375327824e14ba89e4b29cc3ec9a
       selector:
         text:
+
+save_persistent_states:
+  name: Save Persistent States
+  description:
+    Save the persistent states (for entities derived from RestoreEntity) immediately.
+    Maintain the normal periodic saving interval.

homeassistant/const.py

@@ -612,6 +612,7 @@ SERVICE_CLOSE_COVER: Final = "close_cover"
 SERVICE_CLOSE_COVER_TILT: Final = "close_cover_tilt"
 SERVICE_OPEN_COVER: Final = "open_cover"
 SERVICE_OPEN_COVER_TILT: Final = "open_cover_tilt"
+SERVICE_SAVE_PERSISTENT_STATES: Final = "save_persistent_states"
 SERVICE_SET_COVER_POSITION: Final = "set_cover_position"
 SERVICE_SET_COVER_TILT_POSITION: Final = "set_cover_tilt_position"
 SERVICE_STOP_COVER: Final = "stop_cover"

homeassistant/helpers/restore_state.py

@@ -100,6 +100,12 @@ class RestoreStateData:
 
         return cast(RestoreStateData, await load_instance(hass))
 
+    @classmethod
+    async def async_save_persistent_states(cls, hass: HomeAssistant) -> None:
+        """Dump states now."""
+        data = await cls.async_get_instance(hass)
+        await data.async_dump_states()
+
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize the restore state data class."""
         self.hass: HomeAssistant = hass

tests/components/homeassistant/test_init.py

@@ -23,6 +23,7 @@ from homeassistant.const import (
     EVENT_CORE_CONFIG_UPDATE,
     SERVICE_HOMEASSISTANT_RESTART,
     SERVICE_HOMEASSISTANT_STOP,
+    SERVICE_SAVE_PERSISTENT_STATES,
     SERVICE_TOGGLE,
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
@@ -543,3 +544,18 @@ async def test_stop_homeassistant(hass):
         assert not mock_check.called
         await hass.async_block_till_done()
         assert mock_restart.called
+
+
+async def test_save_persistent_states(hass):
+    """Test we can call save_persistent_states."""
+    await async_setup_component(hass, "homeassistant", {})
+    with patch(
+        "homeassistant.helpers.restore_state.RestoreStateData.async_save_persistent_states",
+        return_value=None,
+    ) as mock_save:
+        await hass.services.async_call(
+            "homeassistant",
+            SERVICE_SAVE_PERSISTENT_STATES,
+            blocking=True,
+        )
+        assert mock_save.called

tests/helpers/test_restore_state.py

@@ -98,6 +98,62 @@ async def test_periodic_write(hass):
     assert not mock_write_data.called
 
 
+async def test_save_persistent_states(hass):
+    """Test that we cancel the currently running job, save the data, and verify the perdiodic job continues."""
+    data = await RestoreStateData.async_get_instance(hass)
+    await hass.async_block_till_done()
+    await data.store.async_save([])
+
+    # Emulate a fresh load
+    hass.data[DATA_RESTORE_STATE_TASK] = None
+
+    entity = RestoreEntity()
+    entity.hass = hass
+    entity.entity_id = "input_boolean.b1"
+
+    with patch(
+        "homeassistant.helpers.restore_state.Store.async_save"
+    ) as mock_write_data:
+        await entity.async_get_last_state()
+        await hass.async_block_till_done()
+
+    # Startup Save
+    assert mock_write_data.called
+
+    with patch(
+        "homeassistant.helpers.restore_state.Store.async_save"
+    ) as mock_write_data:
+        async_fire_time_changed(hass, dt_util.utcnow() + timedelta(minutes=10))
+        await hass.async_block_till_done()
+
+    # Not quite the first interval
+    assert not mock_write_data.called
+
+    with patch(
+        "homeassistant.helpers.restore_state.Store.async_save"
+    ) as mock_write_data:
+        await RestoreStateData.async_save_persistent_states(hass)
+        await hass.async_block_till_done()
+
+    assert mock_write_data.called
+
+    with patch(
+        "homeassistant.helpers.restore_state.Store.async_save"
+    ) as mock_write_data:
+        async_fire_time_changed(hass, dt_util.utcnow() + timedelta(minutes=20))
+        await hass.async_block_till_done()
+    # Verify still saving
+    assert mock_write_data.called
+
+    with patch(
+        "homeassistant.helpers.restore_state.Store.async_save"
+    ) as mock_write_data:
+        hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
+        await hass.async_block_till_done()
+    # Verify normal shutdown
+    assert mock_write_data.called
+
+
 async def test_hass_starting(hass):
     """Test that we cache data."""
     hass.state = CoreState.starting