Стабильные тесты в SQLAlchemy 2.0 с SAVEPOINT и pytest

Как перенести фикстуры с SQLAlchemy 1.4 на 2.0: SAVEPOINT, pytest и join_transaction_mode=create_savepoint. Без монкипатча commit→flush, изоляция тестов.

При миграции набора тестов с SQLAlchemy 1.3/1.4 на 2.0 схемы, опиравшиеся на вложенные транзакции через SAVEPOINT в Postgres, могут начать давать нестабильные результаты. Часто используют подход: однажды на модуль подготовить общие фикстуры, а каждый тест изолировать вложенной транзакцией с откатом. В 1.4 это нередко работало с простым монкипатчем, который заменял session.commit на session.flush. После перехода на 2.0 тесты могут сообщать, что не видят общие данные, хотя низкоуровневых ошибок SQLAlchemy нет.

Исходная схема

Шаблон ниже показывает сессию с областью видимости «package», которая подготавливает общие данные, и вложенную транзакцию с областью «function», изолирующую каждый отдельный тест. Монкипатч «commit→flush» гарантирует, что ничего не вырвется за пределы оборачивающей транзакции.

@pytest.fixture(scope="package")
def pkg_sess(suite_engine: Engine, pkg_patch) -> Session:
    link = suite_engine.connect()
    outer_tx = link.begin()
    db_sess = Session(bind=link)
    pkg_patch.setattr(db_sess, "commit", db_sess.flush)
    seed_shared_state(db_sess)
    try:
        yield db_sess
    finally:
        outer_tx.rollback()
        link.close()


@pytest.fixture(scope="function")
def txn_sess(pkg_sess):
    """Создать SAVEPOINT, чтобы изменения конкретного теста можно было откатить, не теряя общие данные"""
    pkg_sess.begin_nested()
    try:
        yield pkg_sess
    finally:
        pkg_sess.rollback()

Что именно ломается

Под SQLAlchemy 2.0 набор начинает падать с сообщениями об отсутствующих общих строках. Подход, который раньше держал общие данные видимыми во вложенных тестах, больше не ведёт себя ожидаемо. После многочисленных проб и ошибок проявилось устойчивое наблюдение: несколько уровней SAVEPOINT воспроизвести как раньше не удаётся. Пакетный уровень вложенности с откатом на уровне теста работает, но более глубокие уровни — например, общие данные на уровне пакета плюс данные на уровне модуля и откат на уровне функции — повторить прежним образом не получилось. В результате один модуль, опиравшийся на тройную вложенность, пришлось немного досевать данными в каждом тесте.

Рабочий подход в SQLAlchemy 2.0

Решение, которое стабильно возвращает исходный замысел: держать явный Connection, оборачивать его во внешнюю транзакцию на уровне пакета или теста по необходимости и указать Session присоединяться через SAVEPOINT. Больше нет нужды монкипатчить commit.

tests/conftest.py

from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session
from sqlalchemy_utils import create_database, database_exists, drop_database

from myapp.db import DeclarativeBase as OrmBase


@pytest.fixture(scope="session")
def qa_engine():
    """Создать отдельную тестовую базу данных, прогнать миграции и удалить её по завершении сессии"""
    eng = create_engine(
        "postgresql+psycopg2://USER:PASSWORD@HOST:PORT/test",
        echo=False,
        future=True,
    )
    if database_exists(eng.url):
        drop_database(eng.url)
    create_database(eng.url)
    OrmBase.metadata.create_all(bind=eng)
    try:
        yield eng
    finally:
        drop_database(eng.url)


@pytest.fixture(scope="function")
def db_session(qa_engine: Engine) -> Session:
    """Открыть транзакцию и использовать SAVEPOINT для изоляции каждого теста"""
    link = qa_engine.connect()
    tx = link.begin()
    try:
        with Session(bind=link, join_transaction_mode="create_savepoint") as db:
            yield db
    finally:
        tx.rollback()
        link.close()

tests/module_a/conftest.py

from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session


@pytest.fixture(scope="package")
def pkg_link(qa_engine: Engine):
    """Один раз на пакет наполнить БД данными, нужными для тестов модуля A.

    Одна и та же связь и внешняя транзакция будут переиспользоваться тестами этого пакета.
    """
    link = qa_engine.connect()
    top_tx = link.begin()
    top_sess = Session(bind=link, join_transaction_mode="create_savepoint")
    # Здесь создайте общие данные для этого пакета

    try:
        yield link
    finally:
        top_tx.rollback()
        link.close()


@pytest.fixture(scope="function")
def db_session(pkg_link):
    """Использовать вложенный SAVEPOINT, чтобы отдельные тесты могли откатываться, не затрагивая общие данные"""
    sp = pkg_link.begin_nested()
    try:
        with Session(bind=pkg_link, join_transaction_mode="create_savepoint") as db:
            yield db
    finally:
        sp.rollback()

Как это решает проблему

Ключевое изменение — отказаться от монкипатча «commit→flush» и явно указать Session, как участвовать в уже начатой транзакции. Режим join_transaction_mode="create_savepoint" заставляет каждый тест работать под SAVEPOINT, пока внешняя транзакция остаётся открытой на уровне пакета. По завершении теста SAVEPOINT откатывается, а предварительно посеянные данные пакета остаются нетронутыми. В этой схеме переопределять commit не нужно; при посеве на уровне пакета можно вызывать commit, и откат SAVEPOINT на уровне теста всё равно отменит только изменения самого теста.

Тесты верхнего уровня ожидают пустую базу данных. Тесты в Module A — предзаполненные данные. Тесты в соседнем Module B — другой набор общих данных. Настройка на уровне пакета позволяет один раз на модуль создать общие данные, сохраняя изоляцию на уровне тестов.

Замеченные ограничения

Несколько вложенных уровней SAVEPOINT не удалось воспроизвести в том же виде под 2.0. Пакетный уровень вложенности с откатом SAVEPOINT на уровне теста работает надёжно. Если нужен ещё один уровень, вероятно, придётся немного подготавливать данные в каждом тесте соответствующего модуля.

Почему это важно

Согласованная изоляция транзакций делает тесты детерминированными и быстрыми. Предпосев общих фикстур один раз на модуль сокращает дублирование и время прогона, а откаты на базе SAVEPOINT удерживают записи каждого теста в его пределах. Адаптация тестовой обвязки к SQLAlchemy 2.0 с явными соединениями и join_transaction_mode возвращает эти свойства без монкипатчинга поведения Session.

Практические выводы

Используйте явный Connection с обрамляющей транзакцией для области, которая владеет общими данными. Создавайте Session каждого теста с join_transaction_mode="create_savepoint", чтобы откаты затрагивали только записи конкретного теста. Откажитесь от монкипатча commit→flush; коммиты во время пакетного посева допустимы, а откат SAVEPOINT на уровне функции оставит общие строки нетронутыми. Если раньше вы полагались на более чем один дополнительный слой SAVEPOINT, будьте готовы слегка досевать данные в каждом тесте таких модулей.

С этими корректировками стратегию тестов с вложенными транзакциями из 1.3/1.4 можно перенести на SQLAlchemy 2.0 с минимальными усилиями и предсказуемым поведением.