diff --git a/homeassistant/helpers/storage.py b/homeassistant/helpers/storage.py index d57a06b74098..b44a250af219 100644 --- a/homeassistant/helpers/storage.py +++ b/homeassistant/helpers/storage.py @@ -28,6 +28,7 @@ from homeassistant.core import ( ) from homeassistant.exceptions import HomeAssistantError, UnsupportedStorageVersionError from homeassistant.util import dt as dt_util, json as json_util +from homeassistant.util.async_ import wait_shared_future from homeassistant.util.file import WriteError, write_utf8_file, write_utf8_file_atomic from homeassistant.util.hass_dict import HassKey @@ -312,7 +313,7 @@ class Store[_T: Mapping[str, Any] | Sequence[Any]]: the second call will wait and return the result of the first call. """ if self._load_future: - return await self._load_future + return await wait_shared_future(self._load_future) self._load_future = self.hass.loop.create_future() try: diff --git a/tests/helpers/test_storage.py b/tests/helpers/test_storage.py index 249dab632fe5..864e6def1c6b 100644 --- a/tests/helpers/test_storage.py +++ b/tests/helpers/test_storage.py @@ -131,6 +131,29 @@ async def test_loading_parallel( assert caplog.text.count(f"Loading data for {store.key}") +async def test_loading_parallel_waiter_cancelled( + hass: HomeAssistant, store: storage.Store +) -> None: + """Verify cancelling a waiting caller does not break the in-progress load.""" + load_started = asyncio.Event() + finish_load = asyncio.Event() + + async def _async_load() -> dict[str, str]: + load_started.set() + await finish_load.wait() + return MOCK_DATA + + with patch.object(store, "_async_load", _async_load): + load_task1 = asyncio.create_task(store.async_load()) + load_task2 = asyncio.create_task(store.async_load()) + await load_started.wait() + load_task2.cancel() + with pytest.raises(asyncio.CancelledError): + await load_task2 + finish_load.set() + assert await load_task1 == MOCK_DATA + + async def test_saving_with_delay( hass: HomeAssistant, store: storage.Store, hass_storage: dict[str, Any] ) -> None: