From 98d5f2fc014a8d2a4771cb79e2e388ef0394fd39 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Tue, 20 Feb 2024 20:01:50 -0600 Subject: [PATCH] Reduce registry overhead in tests (#110955) * Avoid scheduling registry loads as tasks in tests Since we patch out async_load in Store, these will not yield to the event loop so it makes sense to await them instead of creating tasks This reduced my local test run times ~2.5% on average * mock out save as well so we do not schedule tasks to save empty data * tweaks * fix lingering files * another one * too much for one PR, reduce * fix targets --- tests/common.py | 50 ++++++++++++++++++++++++++--------- tests/helpers/test_storage.py | 15 ++++++----- 2 files changed, 46 insertions(+), 19 deletions(-) diff --git a/tests/common.py b/tests/common.py index e7648bba303..3ed865af00b 100644 --- a/tests/common.py +++ b/tests/common.py @@ -17,7 +17,7 @@ import pathlib import threading import time from types import ModuleType -from typing import Any, NoReturn +from typing import Any, NoReturn, TypeVar from unittest.mock import AsyncMock, Mock, patch from aiohttp.test_utils import unused_port as get_test_instance_port # noqa: F401 @@ -194,13 +194,26 @@ def get_test_home_assistant() -> Generator[HomeAssistant, None, None]: loop.close() +_T = TypeVar("_T", bound=Mapping[str, Any] | Sequence[Any]) + + +class StoreWithoutWriteLoad(storage.Store[_T]): + """Fake store that does not write or load. Used for testing.""" + + async def async_save(self, *args: Any, **kwargs: Any) -> None: + """Save the data.""" + + @asynccontextmanager async def async_test_home_assistant( event_loop: asyncio.AbstractEventLoop | None = None, load_registries: bool = True, + storage_dir: str | None = None, ) -> AsyncGenerator[HomeAssistant, None]: """Return a Home Assistant object pointing at test config dir.""" hass = HomeAssistant(get_test_config_dir()) + if storage_dir: + hass.config.config_dir = storage_dir store = auth_store.AuthStore(hass) hass.auth = auth.AuthManager(hass, store, {}, {}) ensure_auth_manager_loaded(hass.auth) @@ -284,23 +297,36 @@ async def async_test_home_assistant( hass ) if load_registries: - with patch( - "homeassistant.helpers.storage.Store.async_load", return_value=None + with patch.object( + StoreWithoutWriteLoad, "async_load", return_value=None + ), patch( + "homeassistant.helpers.area_registry.AreaRegistryStore", + StoreWithoutWriteLoad, + ), patch( + "homeassistant.helpers.device_registry.DeviceRegistryStore", + StoreWithoutWriteLoad, + ), patch( + "homeassistant.helpers.entity_registry.EntityRegistryStore", + StoreWithoutWriteLoad, + ), patch( + "homeassistant.helpers.storage.Store", # Floor & label registry are different + StoreWithoutWriteLoad, + ), patch( + "homeassistant.helpers.issue_registry.IssueRegistryStore", + StoreWithoutWriteLoad, ), patch( "homeassistant.helpers.restore_state.RestoreStateData.async_setup_dump", return_value=None, ), patch( "homeassistant.helpers.restore_state.start.async_at_start", ): - await asyncio.gather( - ar.async_load(hass), - dr.async_load(hass), - er.async_load(hass), - fr.async_load(hass), - ir.async_load(hass), - lr.async_load(hass), - rs.async_load(hass), - ) + await ar.async_load(hass) + await dr.async_load(hass) + await er.async_load(hass) + await fr.async_load(hass) + await ir.async_load(hass) + await lr.async_load(hass) + await rs.async_load(hass) hass.data[bootstrap.DATA_REGISTRIES_LOADED] = None hass.set_state(CoreState.running) diff --git a/tests/helpers/test_storage.py b/tests/helpers/test_storage.py index 92f6c3c18a5..ab1889eccea 100644 --- a/tests/helpers/test_storage.py +++ b/tests/helpers/test_storage.py @@ -560,10 +560,10 @@ async def test_loading_corrupt_core_file( tmpdir: py.path.local, caplog: pytest.LogCaptureFixture ) -> None: """Test we handle unrecoverable corruption in a core file.""" - async with async_test_home_assistant() as hass: - tmp_storage = await hass.async_add_executor_job(tmpdir.mkdir, "temp_storage") - hass.config.config_dir = tmp_storage + loop = asyncio.get_running_loop() + tmp_storage = await loop.run_in_executor(None, tmpdir.mkdir, "temp_storage") + async with async_test_home_assistant(storage_dir=tmp_storage) as hass: storage_key = "core.anything" store = storage.Store( hass, MOCK_VERSION_2, storage_key, minor_version=MOCK_MINOR_VERSION_1 @@ -618,13 +618,14 @@ async def test_loading_corrupt_file_known_domain( tmpdir: py.path.local, caplog: pytest.LogCaptureFixture ) -> None: """Test we handle unrecoverable corruption for a known domain.""" - async with async_test_home_assistant() as hass: + + loop = asyncio.get_running_loop() + tmp_storage = await loop.run_in_executor(None, tmpdir.mkdir, "temp_storage") + + async with async_test_home_assistant(storage_dir=tmp_storage) as hass: hass.config.components.add("testdomain") storage_key = "testdomain.testkey" - tmp_storage = await hass.async_add_executor_job(tmpdir.mkdir, "temp_storage") - hass.config.config_dir = tmp_storage - store = storage.Store( hass, MOCK_VERSION_2, storage_key, minor_version=MOCK_MINOR_VERSION_1 )