From 3154c3c962d5eb246047ddb84567ec6bf938bbaa Mon Sep 17 00:00:00 2001 From: Artur Pragacz <49985303+arturpragacz@users.noreply.github.com> Date: Sun, 8 Mar 2026 10:39:53 +0100 Subject: [PATCH] Make restore state resilient to extra_restore_state_data errors (#165086) --- homeassistant/helpers/restore_state.py | 45 +++++++++--- tests/helpers/test_restore_state.py | 99 +++++++++++++++++++++++++- 2 files changed, 130 insertions(+), 14 deletions(-) diff --git a/homeassistant/helpers/restore_state.py b/homeassistant/helpers/restore_state.py index 59f802e2448..81e9d7ed68e 100644 --- a/homeassistant/helpers/restore_state.py +++ b/homeassistant/helpers/restore_state.py @@ -181,15 +181,24 @@ class RestoreStateData: } # Start with the currently registered states - stored_states = [ - StoredState( - current_states_by_entity_id[entity_id], - entity.extra_restore_state_data, - now, + stored_states: list[StoredState] = [] + for entity_id, entity in self.entities.items(): + if entity_id not in current_states_by_entity_id: + continue + try: + extra_data = entity.extra_restore_state_data + except Exception: + _LOGGER.exception( + "Error getting extra restore state data for %s", entity_id + ) + continue + stored_states.append( + StoredState( + current_states_by_entity_id[entity_id], + extra_data, + now, + ) ) - for entity_id, entity in self.entities.items() - if entity_id in current_states_by_entity_id - ] expiration_time = now - STATE_EXPIRATION for entity_id, stored_state in self.last_states.items(): @@ -219,6 +228,8 @@ class RestoreStateData: ) except HomeAssistantError as exc: _LOGGER.error("Error saving current states", exc_info=exc) + except Exception: + _LOGGER.exception("Unexpected error saving current states") @callback def async_setup_dump(self, *args: Any) -> None: @@ -258,13 +269,15 @@ class RestoreStateData: @callback def async_restore_entity_removed( - self, entity_id: str, extra_data: ExtraStoredData | None + self, + entity_id: str, + state: State | None, + extra_data: ExtraStoredData | None, ) -> None: """Unregister this entity from saving state.""" # When an entity is being removed from hass, store its last state. This # allows us to support state restoration if the entity is removed, then # re-added while hass is still running. - state = self.hass.states.get(entity_id) # To fully mimic all the attribute data types when loaded from storage, # we're going to serialize it to JSON and then re-load it. if state is not None: @@ -287,8 +300,18 @@ class RestoreEntity(Entity): async def async_internal_will_remove_from_hass(self) -> None: """Run when entity will be removed from hass.""" + try: + extra_data = self.extra_restore_state_data + except Exception: + _LOGGER.exception( + "Error getting extra restore state data for %s", self.entity_id + ) + state = None + extra_data = None + else: + state = self.hass.states.get(self.entity_id) async_get(self.hass).async_restore_entity_removed( - self.entity_id, self.extra_restore_state_data + self.entity_id, state, extra_data ) await super().async_internal_will_remove_from_hass() diff --git a/tests/helpers/test_restore_state.py b/tests/helpers/test_restore_state.py index 7adb3dd5b5e..6320858a2a4 100644 --- a/tests/helpers/test_restore_state.py +++ b/tests/helpers/test_restore_state.py @@ -6,6 +6,8 @@ import logging from typing import Any from unittest.mock import Mock, patch +import pytest + from homeassistant.const import EVENT_HOMEASSISTANT_START, EVENT_HOMEASSISTANT_STOP from homeassistant.core import CoreState, HomeAssistant, State from homeassistant.exceptions import HomeAssistantError @@ -16,6 +18,7 @@ from homeassistant.helpers.reload import async_get_platform_without_config_entry from homeassistant.helpers.restore_state import ( DATA_RESTORE_STATE, STORAGE_KEY, + ExtraStoredData, RestoreEntity, RestoreStateData, StoredState, @@ -342,8 +345,12 @@ async def test_dump_data(hass: HomeAssistant) -> None: assert state1["state"]["state"] == "off" -async def test_dump_error(hass: HomeAssistant) -> None: - """Test that we cache data.""" +@pytest.mark.parametrize( + "exception", + [HomeAssistantError, RuntimeError], +) +async def test_dump_error(hass: HomeAssistant, exception: type[Exception]) -> None: + """Test that errors during save are caught.""" states = [ State("input_boolean.b0", "on"), State("input_boolean.b1", "on"), @@ -368,7 +375,7 @@ async def test_dump_error(hass: HomeAssistant) -> None: with patch( "homeassistant.helpers.restore_state.Store.async_save", - side_effect=HomeAssistantError, + side_effect=exception, ) as mock_write_data: await data.async_dump_states() @@ -534,3 +541,89 @@ async def test_restore_entity_end_to_end( assert len(storage_data) == 1 assert storage_data[0]["state"]["entity_id"] == entity_id assert storage_data[0]["state"]["state"] == "stored" + + +async def test_dump_states_with_failing_extra_data( + hass: HomeAssistant, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test that a failing extra_restore_state_data skips only that entity.""" + + class BadRestoreEntity(RestoreEntity): + """Entity that raises on extra_restore_state_data.""" + + @property + def extra_restore_state_data(self) -> ExtraStoredData | None: + raise RuntimeError("Unexpected error") + + states = [ + State("input_boolean.good", "on"), + State("input_boolean.bad", "on"), + ] + + platform = MockEntityPlatform(hass, domain="input_boolean") + + good_entity = RestoreEntity() + good_entity.hass = hass + good_entity.entity_id = "input_boolean.good" + await platform.async_add_entities([good_entity]) + + bad_entity = BadRestoreEntity() + bad_entity.hass = hass + bad_entity.entity_id = "input_boolean.bad" + await platform.async_add_entities([bad_entity]) + + for state in states: + hass.states.async_set(state.entity_id, state.state, state.attributes) + + data = async_get(hass) + + with patch( + "homeassistant.helpers.restore_state.Store.async_save" + ) as mock_write_data: + await data.async_dump_states() + + assert mock_write_data.called + written_states = mock_write_data.mock_calls[0][1][0] + + # Only the good entity should be saved + assert len(written_states) == 1 + state0 = json_round_trip(written_states[0]) + assert state0["state"]["entity_id"] == "input_boolean.good" + assert state0["state"]["state"] == "on" + + assert "Error getting extra restore state data for input_boolean.bad" in caplog.text + + +async def test_entity_removal_with_failing_extra_data( + hass: HomeAssistant, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test that entity removal succeeds even if extra_restore_state_data raises.""" + + class BadRestoreEntity(RestoreEntity): + """Entity that raises on extra_restore_state_data.""" + + @property + def extra_restore_state_data(self) -> ExtraStoredData | None: + raise RuntimeError("Unexpected error") + + platform = MockEntityPlatform(hass, domain="input_boolean") + entity = BadRestoreEntity() + entity.hass = hass + entity.entity_id = "input_boolean.bad" + await platform.async_add_entities([entity]) + + hass.states.async_set("input_boolean.bad", "on") + + data = async_get(hass) + assert "input_boolean.bad" in data.entities + + await entity.async_remove() + + # Entity should be unregistered + assert "input_boolean.bad" not in data.entities + # No last state should be saved since extra data failed + assert "input_boolean.bad" not in data.last_states + + assert "Error getting extra restore state data for input_boolean.bad" in caplog.text