Make the async_setup_recorder_instance fixture a context manager (#121086)

This commit is contained in:
Erik Montnemery 2024-07-03 21:05:34 +02:00 committed by GitHub
parent 61f1c8d963
commit 5029da6919
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 51 additions and 22 deletions

View file

@ -100,7 +100,7 @@ _TEST_FIXTURES: dict[str, list[str] | str] = {
"aiohttp_client": "ClientSessionGenerator", "aiohttp_client": "ClientSessionGenerator",
"aiohttp_server": "Callable[[], TestServer]", "aiohttp_server": "Callable[[], TestServer]",
"area_registry": "AreaRegistry", "area_registry": "AreaRegistry",
"async_setup_recorder_instance": "RecorderInstanceGenerator", "async_test_recorder": "RecorderInstanceGenerator",
"caplog": "pytest.LogCaptureFixture", "caplog": "pytest.LogCaptureFixture",
"capsys": "pytest.CaptureFixture[str]", "capsys": "pytest.CaptureFixture[str]",
"current_request_with_host": "None", "current_request_with_host": "None",

View file

@ -13,7 +13,7 @@ from tests.typing import RecorderInstanceGenerator
@pytest.fixture @pytest.fixture
async def mock_recorder_before_hass( async def mock_recorder_before_hass(
async_setup_recorder_instance: RecorderInstanceGenerator, async_test_recorder: RecorderInstanceGenerator,
) -> None: ) -> None:
"""Set up recorder.""" """Set up recorder."""

View file

@ -29,7 +29,7 @@ from tests.typing import RecorderInstanceGenerator
@pytest.fixture @pytest.fixture
async def mock_recorder_before_hass( async def mock_recorder_before_hass(
async_setup_recorder_instance: RecorderInstanceGenerator, async_test_recorder: RecorderInstanceGenerator,
) -> None: ) -> None:
"""Set up recorder.""" """Set up recorder."""

View file

@ -40,7 +40,7 @@ def _count_entity_id_in_states_meta(
@pytest.fixture @pytest.fixture
async def mock_recorder_before_hass( async def mock_recorder_before_hass(
async_setup_recorder_instance: RecorderInstanceGenerator, async_test_recorder: RecorderInstanceGenerator,
) -> None: ) -> None:
"""Set up recorder.""" """Set up recorder."""

View file

@ -47,7 +47,7 @@ from tests.typing import RecorderInstanceGenerator
@pytest.fixture @pytest.fixture
async def mock_recorder_before_hass( async def mock_recorder_before_hass(
async_setup_recorder_instance: RecorderInstanceGenerator, async_test_recorder: RecorderInstanceGenerator,
) -> None: ) -> None:
"""Set up recorder.""" """Set up recorder."""

View file

@ -33,7 +33,7 @@ from tests.typing import RecorderInstanceGenerator
@pytest.fixture @pytest.fixture
async def mock_recorder_before_hass( async def mock_recorder_before_hass(
async_setup_recorder_instance: RecorderInstanceGenerator, async_test_recorder: RecorderInstanceGenerator,
) -> None: ) -> None:
"""Set up recorder.""" """Set up recorder."""

View file

@ -33,7 +33,7 @@ from tests.typing import RecorderInstanceGenerator
@pytest.fixture @pytest.fixture
async def mock_recorder_before_hass( async def mock_recorder_before_hass(
async_setup_recorder_instance: RecorderInstanceGenerator, async_test_recorder: RecorderInstanceGenerator,
) -> None: ) -> None:
"""Set up recorder.""" """Set up recorder."""

View file

@ -42,7 +42,7 @@ from tests.typing import RecorderInstanceGenerator
@pytest.fixture @pytest.fixture
async def mock_recorder_before_hass( async def mock_recorder_before_hass(
async_setup_recorder_instance: RecorderInstanceGenerator, async_test_recorder: RecorderInstanceGenerator,
) -> None: ) -> None:
"""Set up recorder.""" """Set up recorder."""

View file

@ -104,7 +104,7 @@ from tests.typing import RecorderInstanceGenerator
@pytest.fixture @pytest.fixture
async def mock_recorder_before_hass( async def mock_recorder_before_hass(
async_setup_recorder_instance: RecorderInstanceGenerator, async_test_recorder: RecorderInstanceGenerator,
) -> None: ) -> None:
"""Set up recorder.""" """Set up recorder."""

View file

@ -40,7 +40,7 @@ from tests.typing import RecorderInstanceGenerator
@pytest.fixture @pytest.fixture
async def mock_recorder_before_hass( async def mock_recorder_before_hass(
async_setup_recorder_instance: RecorderInstanceGenerator, async_test_recorder: RecorderInstanceGenerator,
) -> None: ) -> None:
"""Set up recorder.""" """Set up recorder."""

View file

@ -51,7 +51,7 @@ from tests.typing import RecorderInstanceGenerator, WebSocketGenerator
@pytest.fixture @pytest.fixture
async def mock_recorder_before_hass( async def mock_recorder_before_hass(
async_setup_recorder_instance: RecorderInstanceGenerator, async_test_recorder: RecorderInstanceGenerator,
) -> None: ) -> None:
"""Set up recorder.""" """Set up recorder."""

View file

@ -48,7 +48,7 @@ from tests.typing import RecorderInstanceGenerator
@pytest.fixture @pytest.fixture
async def mock_recorder_before_hass( async def mock_recorder_before_hass(
async_setup_recorder_instance: RecorderInstanceGenerator, async_test_recorder: RecorderInstanceGenerator,
) -> None: ) -> None:
"""Set up recorder.""" """Set up recorder."""

View file

@ -95,7 +95,7 @@ KW_SENSOR_ATTRIBUTES = {
@pytest.fixture @pytest.fixture
async def mock_recorder_before_hass( async def mock_recorder_before_hass(
async_setup_recorder_instance: RecorderInstanceGenerator, async_test_recorder: RecorderInstanceGenerator,
) -> None: ) -> None:
"""Set up recorder patches.""" """Set up recorder patches."""

View file

@ -4,7 +4,7 @@ from __future__ import annotations
import asyncio import asyncio
from collections.abc import AsyncGenerator, Callable, Coroutine, Generator from collections.abc import AsyncGenerator, Callable, Coroutine, Generator
from contextlib import asynccontextmanager, contextmanager from contextlib import AsyncExitStack, asynccontextmanager, contextmanager
import datetime import datetime
import functools import functools
import gc import gc
@ -1385,7 +1385,7 @@ async def _async_init_recorder_component(
@pytest.fixture @pytest.fixture
async def async_setup_recorder_instance( async def async_test_recorder(
recorder_db_url: str, recorder_db_url: str,
enable_nightly_purge: bool, enable_nightly_purge: bool,
enable_statistics: bool, enable_statistics: bool,
@ -1394,7 +1394,7 @@ async def async_setup_recorder_instance(
enable_migrate_event_type_ids: bool, enable_migrate_event_type_ids: bool,
enable_migrate_entity_ids: bool, enable_migrate_entity_ids: bool,
) -> AsyncGenerator[RecorderInstanceGenerator]: ) -> AsyncGenerator[RecorderInstanceGenerator]:
"""Yield callable to setup recorder instance.""" """Yield context manager to setup recorder instance."""
# pylint: disable-next=import-outside-toplevel # pylint: disable-next=import-outside-toplevel
from homeassistant.components import recorder from homeassistant.components import recorder
@ -1477,12 +1477,13 @@ async def async_setup_recorder_instance(
), ),
): ):
async def async_setup_recorder( @asynccontextmanager
async def async_test_recorder(
hass: HomeAssistant, hass: HomeAssistant,
config: ConfigType | None = None, config: ConfigType | None = None,
*, *,
wait_recorder: bool = True, wait_recorder: bool = True,
) -> recorder.Recorder: ) -> AsyncGenerator[recorder.Recorder]:
"""Setup and return recorder instance.""" # noqa: D401 """Setup and return recorder instance.""" # noqa: D401
await _async_init_recorder_component(hass, config, recorder_db_url) await _async_init_recorder_component(hass, config, recorder_db_url)
await hass.async_block_till_done() await hass.async_block_till_done()
@ -1490,7 +1491,34 @@ async def async_setup_recorder_instance(
# The recorder's worker is not started until Home Assistant is running # The recorder's worker is not started until Home Assistant is running
if hass.state is CoreState.running and wait_recorder: if hass.state is CoreState.running and wait_recorder:
await async_recorder_block_till_done(hass) await async_recorder_block_till_done(hass)
return instance try:
yield instance
finally:
if instance.is_alive():
await instance._async_shutdown(None)
yield async_test_recorder
@pytest.fixture
async def async_setup_recorder_instance(
async_test_recorder: RecorderInstanceGenerator,
) -> AsyncGenerator[RecorderInstanceGenerator]:
"""Yield callable to setup recorder instance."""
async with AsyncExitStack() as stack:
async def async_setup_recorder(
hass: HomeAssistant,
config: ConfigType | None = None,
*,
wait_recorder: bool = True,
) -> AsyncGenerator[recorder.Recorder]:
"""Set up and return recorder instance."""
return await stack.enter_async_context(
async_test_recorder(hass, config, wait_recorder=wait_recorder)
)
yield async_setup_recorder yield async_setup_recorder
@ -1498,11 +1526,12 @@ async def async_setup_recorder_instance(
@pytest.fixture @pytest.fixture
async def recorder_mock( async def recorder_mock(
recorder_config: dict[str, Any] | None, recorder_config: dict[str, Any] | None,
async_setup_recorder_instance: RecorderInstanceGenerator, async_test_recorder: RecorderInstanceGenerator,
hass: HomeAssistant, hass: HomeAssistant,
) -> recorder.Recorder: ) -> AsyncGenerator[recorder.Recorder]:
"""Fixture with in-memory recorder.""" """Fixture with in-memory recorder."""
return await async_setup_recorder_instance(hass, recorder_config) async with async_test_recorder(hass, recorder_config) as instance:
yield instance
@pytest.fixture @pytest.fixture