Add opt-in to Store for serializing in an executor (#157263)

tests/helpers/test_storage.py

@@ -27,7 +27,7 @@ from homeassistant.core import (
 )
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import issue_registry as ir, storage
-from homeassistant.helpers.json import json_bytes
+from homeassistant.helpers.json import json_bytes, prepare_save_json
 from homeassistant.util import dt as dt_util
 from homeassistant.util.color import RGBColor
 
@@ -177,17 +177,41 @@ async def test_saving_with_delay_threading(tmp_path: Path) -> None:
             calls.append("callback")
             return MOCK_DATA2
 
-        store = storage.Store(hass, MOCK_VERSION, MOCK_KEY)
-        store.async_delay_save(data_producer_thread_safe, 1)
+        def mock_prepare_thread_safe(*args, **kwargs):
+            """Mock prepare thread safe."""
+            assert threading.get_ident() != hass.loop_thread_id
+            return prepare_save_json(*args, **kwargs)
 
-        async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=1))
-        await hass.async_block_till_done()
+        def mock_prepare_not_thread_safe(*args, **kwargs):
+            """Mock prepare not thread safe."""
+            assert threading.get_ident() == hass.loop_thread_id
+            return prepare_save_json(*args, **kwargs)
 
-        store = storage.Store(hass, MOCK_VERSION, MOCK_KEY2)
-        store.async_delay_save(data_producer_callback, 1)
+        with patch(
+            "homeassistant.helpers.storage.json_helper.prepare_save_json",
+            wraps=mock_prepare_thread_safe,
+        ) as mock_prepare:
+            store = storage.Store(
+                hass, MOCK_VERSION, MOCK_KEY, serialize_in_event_loop=False
+            )
+            store.async_delay_save(data_producer_thread_safe, 1)
 
-        async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=1))
-        await hass.async_block_till_done()
+            async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=1))
+            await hass.async_block_till_done()
+
+            mock_prepare.assert_called_once()
+
+        with patch(
+            "homeassistant.helpers.storage.json_helper.prepare_save_json",
+            wraps=mock_prepare_not_thread_safe,
+        ) as mock_prepare:
+            store = storage.Store(hass, MOCK_VERSION, MOCK_KEY2)
+            store.async_delay_save(data_producer_callback, 1)
+
+            async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=1))
+            await hass.async_block_till_done()
+
+            mock_prepare.assert_called_once()
 
         assert calls == ["thread_safe", "callback"]
         expected_data = (
@@ -216,6 +240,75 @@ async def test_saving_with_delay_threading(tmp_path: Path) -> None:
         await hass.async_stop(force=True)
 
 
+async def test_saving_with_threading(tmp_path: Path) -> None:
+    """Test thread handling when saving."""
+
+    async def assert_storage_data(store_key: str, expected_data: str) -> None:
+        """Assert storage data."""
+
+        def read_storage_data(store_key: str) -> str:
+            """Read storage data."""
+            return Path(tmp_path / f".storage/{store_key}").read_text(encoding="utf-8")
+
+        store_data = await asyncio.to_thread(read_storage_data, store_key)
+        assert store_data == expected_data
+
+    async with async_test_home_assistant(config_dir=tmp_path) as hass:
+
+        def mock_prepare_thread_safe(*args, **kwargs):
+            """Mock prepare thread safe."""
+            assert threading.get_ident() != hass.loop_thread_id
+            return prepare_save_json(*args, **kwargs)
+
+        def mock_prepare_not_thread_safe(*args, **kwargs):
+            """Mock prepare not thread safe."""
+            assert threading.get_ident() == hass.loop_thread_id
+            return prepare_save_json(*args, **kwargs)
+
+        with patch(
+            "homeassistant.helpers.storage.json_helper.prepare_save_json",
+            wraps=mock_prepare_thread_safe,
+        ) as mock_prepare:
+            store = storage.Store(
+                hass, MOCK_VERSION, MOCK_KEY, serialize_in_event_loop=False
+            )
+            await store.async_save(MOCK_DATA)
+            mock_prepare.assert_called_once()
+
+        with patch(
+            "homeassistant.helpers.storage.json_helper.prepare_save_json",
+            wraps=mock_prepare_not_thread_safe,
+        ) as mock_prepare:
+            store = storage.Store(hass, MOCK_VERSION, MOCK_KEY2)
+            await store.async_save(MOCK_DATA2)
+            mock_prepare.assert_called_once()
+
+        expected_data = (
+            "{\n"
+            '  "version": 1,\n'
+            '  "minor_version": 1,\n'
+            '  "key": "storage-test",\n'
+            '  "data": {\n'
+            '    "hello": "world"\n'
+            "  }\n"
+            "}"
+        )
+        await assert_storage_data(MOCK_KEY, expected_data)
+        expected_data = (
+            "{\n"
+            '  "version": 1,\n'
+            '  "minor_version": 1,\n'
+            '  "key": "storage-test-2",\n'
+            '  "data": {\n'
+            '    "goodbye": "cruel world"\n'
+            "  }\n"
+            "}"
+        )
+        await assert_storage_data(MOCK_KEY2, expected_data)
+
+        await hass.async_stop(force=True)
+
+
 async def test_saving_with_delay_churn_reduction(
     hass: HomeAssistant,
     store: storage.Store,