127 lines
4.4 KiB
Python
127 lines
4.4 KiB
Python
"""Fixtures for the recorder component tests."""
|
|
|
|
from collections.abc import AsyncGenerator, Generator
|
|
from dataclasses import dataclass
|
|
from functools import partial
|
|
import threading
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
|
|
from homeassistant.components import recorder
|
|
from homeassistant.components.recorder import db_schema
|
|
from homeassistant.components.recorder.util import session_scope
|
|
from homeassistant.core import HomeAssistant
|
|
|
|
|
|
def pytest_configure(config):
|
|
"""Add custom skip_on_db_engine marker."""
|
|
config.addinivalue_line(
|
|
"markers",
|
|
"skip_on_db_engine(engine): mark test to run only on named DB engine(s)",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def skip_by_db_engine(request: pytest.FixtureRequest, recorder_db_url: str) -> None:
|
|
"""Fixture to skip tests on unsupported DB engines.
|
|
|
|
Mark the test with @pytest.mark.skip_on_db_engine("mysql") to skip on mysql, or
|
|
@pytest.mark.skip_on_db_engine(["mysql", "sqlite"]) to skip on mysql and sqlite.
|
|
"""
|
|
if request.node.get_closest_marker("skip_on_db_engine"):
|
|
skip_on_db_engine = request.node.get_closest_marker("skip_on_db_engine").args[0]
|
|
if isinstance(skip_on_db_engine, str):
|
|
skip_on_db_engine = [skip_on_db_engine]
|
|
db_engine = recorder_db_url.partition("://")[0]
|
|
if db_engine in skip_on_db_engine:
|
|
pytest.skip(f"skipped for DB engine: {db_engine}")
|
|
|
|
|
|
@pytest.fixture
|
|
def recorder_dialect_name(hass: HomeAssistant, db_engine: str) -> Generator[None]:
|
|
"""Patch the recorder dialect."""
|
|
if instance := hass.data.get(recorder.DATA_INSTANCE):
|
|
instance.__dict__.pop("dialect_name", None)
|
|
with patch.object(instance, "_dialect_name", db_engine):
|
|
yield
|
|
instance.__dict__.pop("dialect_name", None)
|
|
else:
|
|
with patch(
|
|
"homeassistant.components.recorder.Recorder.dialect_name", db_engine
|
|
):
|
|
yield
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class InstrumentedMigration:
|
|
"""Container to aid controlling migration progress."""
|
|
|
|
migration_done: threading.Event
|
|
migration_stall: threading.Event
|
|
migration_started: threading.Event
|
|
migration_version: int | None
|
|
apply_update_mock: Mock
|
|
|
|
|
|
@pytest.fixture
|
|
async def instrument_migration(
|
|
hass: HomeAssistant,
|
|
) -> AsyncGenerator[InstrumentedMigration]:
|
|
"""Instrument recorder migration."""
|
|
|
|
real_migrate_schema_live = recorder.migration.migrate_schema_live
|
|
real_migrate_schema_non_live = recorder.migration.migrate_schema_non_live
|
|
real_apply_update = recorder.migration._apply_update
|
|
|
|
def _instrument_migrate_schema(real_func, *args):
|
|
"""Control migration progress and check results."""
|
|
instrumented_migration.migration_started.set()
|
|
|
|
try:
|
|
migration_result = real_func(*args)
|
|
except Exception:
|
|
instrumented_migration.migration_done.set()
|
|
raise
|
|
|
|
# Check and report the outcome of the migration; if migration fails
|
|
# the recorder will silently create a new database.
|
|
with session_scope(hass=hass, read_only=True) as session:
|
|
res = (
|
|
session.query(db_schema.SchemaChanges)
|
|
.order_by(db_schema.SchemaChanges.change_id.desc())
|
|
.first()
|
|
)
|
|
instrumented_migration.migration_version = res.schema_version
|
|
instrumented_migration.migration_done.set()
|
|
return migration_result
|
|
|
|
def _instrument_apply_update(*args):
|
|
"""Control migration progress."""
|
|
instrumented_migration.migration_stall.wait()
|
|
real_apply_update(*args)
|
|
|
|
with (
|
|
patch(
|
|
"homeassistant.components.recorder.migration.migrate_schema_live",
|
|
wraps=partial(_instrument_migrate_schema, real_migrate_schema_live),
|
|
),
|
|
patch(
|
|
"homeassistant.components.recorder.migration.migrate_schema_non_live",
|
|
wraps=partial(_instrument_migrate_schema, real_migrate_schema_non_live),
|
|
),
|
|
patch(
|
|
"homeassistant.components.recorder.migration._apply_update",
|
|
wraps=_instrument_apply_update,
|
|
) as apply_update_mock,
|
|
):
|
|
instrumented_migration = InstrumentedMigration(
|
|
migration_done=threading.Event(),
|
|
migration_stall=threading.Event(),
|
|
migration_started=threading.Event(),
|
|
migration_version=None,
|
|
apply_update_mock=apply_update_mock,
|
|
)
|
|
|
|
yield instrumented_migration
|