Created
December 2, 2025 23:59
-
-
Save a1d4r/d469e1eaac6e3a2bf6a60452c5dcad38 to your computer and use it in GitHub Desktop.
Setup for fast database tests
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import contextlib | |
| from collections.abc import AsyncIterator | |
| import psycopg | |
| import pytest | |
| import pytest_alembic | |
| from filelock import FileLock | |
| from pytest_alembic.config import Config | |
| from pytest_postgresql.janitor import DatabaseJanitor | |
| from sqlalchemy import NullPool | |
| from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine | |
| from app.core.logger import logger | |
| from app.core.settings import DatabaseSettings, db_settings | |
| def _create_janitor( | |
| settings: DatabaseSettings, template_dbname: str | None = None | |
| ) -> DatabaseJanitor: | |
| """Создать DatabaseJanitor для управления БД.""" | |
| return DatabaseJanitor( | |
| user=settings.username.get_secret_value(), | |
| host=settings.host, | |
| port=settings.port, | |
| dbname=settings.name, | |
| template_dbname=template_dbname, | |
| version="16.0", | |
| password=settings.password.get_secret_value(), | |
| ) | |
| @pytest.fixture | |
| def test_database_settings(worker_id: str) -> DatabaseSettings: | |
| """Create database settings with respect to xdist.""" | |
| xdist_db_settings = db_settings.model_copy() | |
| xdist_db_settings.name = f"{xdist_db_settings.name}_{worker_id}" | |
| return xdist_db_settings | |
| @pytest.fixture(scope="session") | |
| def template_database_settings() -> DatabaseSettings: | |
| """Настройки для шаблона БД.""" | |
| settings = db_settings.model_copy() | |
| settings.name = f"{settings.name}_template" | |
| return settings | |
| @pytest.fixture(scope="session") | |
| def _template_db( | |
| tmp_path_factory: pytest.TempPathFactory, | |
| worker_id: str, | |
| template_database_settings: DatabaseSettings, | |
| ): | |
| """Создать шаблон БД (один раз за сессию).""" | |
| janitor = _create_janitor(template_database_settings) | |
| def create_template_database(): | |
| with contextlib.suppress(psycopg.errors.DatabaseError): | |
| janitor.drop() | |
| janitor.init() | |
| # Применяем миграции | |
| engine = create_async_engine(template_database_settings.url, poolclass=NullPool) | |
| config = Config.from_raw_config({}) | |
| with pytest_alembic.runner(config=config, engine=engine) as runner: | |
| runner.migrate_up_to("heads", return_current=False) | |
| logger.info(f"Template database '{template_database_settings.name}' created") | |
| if worker_id == "master": | |
| create_template_database() | |
| yield template_database_settings.name | |
| with contextlib.suppress(psycopg.errors.DatabaseError): | |
| janitor.drop() | |
| else: | |
| root_tmp_dir = tmp_path_factory.getbasetemp().parent | |
| fn = root_tmp_dir / "template_db_created" | |
| with FileLock(str(fn) + ".lock"): | |
| if not fn.is_file(): | |
| create_template_database() | |
| fn.write_text(template_database_settings.name) | |
| yield template_database_settings.name | |
| @pytest.fixture | |
| def _empty_postgres(test_database_settings: DatabaseSettings): | |
| """Пустая БД для тестов миграций pytest-alembic.""" | |
| janitor = _create_janitor(test_database_settings) | |
| with contextlib.suppress(psycopg.errors.DatabaseError): | |
| janitor.drop() | |
| janitor.init() | |
| @pytest.fixture | |
| def _postgres(_template_db: str, test_database_settings: DatabaseSettings): | |
| """БД из template для обычных тестов.""" | |
| if not test_database_settings.name.startswith("test"): | |
| raise RuntimeError("Running tests on non-test database is forbidden") | |
| janitor = _create_janitor(test_database_settings, template_dbname=_template_db) | |
| with contextlib.suppress(psycopg.errors.DatabaseError): | |
| janitor.drop() | |
| janitor.init() | |
| yield | |
| with contextlib.suppress(psycopg.errors.DatabaseError): | |
| janitor.drop() | |
| @pytest.fixture | |
| def alembic_engine(_empty_postgres, test_database_settings: DatabaseSettings) -> AsyncEngine: | |
| """Engine для тестов миграций pytest-alembic.""" | |
| return create_async_engine(test_database_settings.url, poolclass=NullPool) | |
| @pytest.fixture | |
| def async_engine(_postgres, test_database_settings: DatabaseSettings) -> AsyncEngine: | |
| """Engine для обычных тестов.""" | |
| return create_async_engine(test_database_settings.url, poolclass=NullPool) | |
| @pytest.fixture | |
| async def session(async_engine: AsyncEngine) -> AsyncIterator[AsyncSession]: | |
| """Session for tests.""" | |
| async with AsyncSession( | |
| async_engine, expire_on_commit=False, autoflush=False, autocommit=False | |
| ) as session: | |
| yield session | |
| @pytest.fixture | |
| async def app_session(async_engine: AsyncEngine) -> AsyncIterator[AsyncSession]: | |
| """Use different session for app itself.""" | |
| async with AsyncSession( | |
| async_engine, expire_on_commit=False, autoflush=False, autocommit=False | |
| ) as session: | |
| yield session |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment