How to test rollbacks in SQLAlchemy

Question

I have a problem with test isolation when testing a logic, that involves a transaction rollback in SQLAlchemy.

Model:

class Product(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    company = db.Column(db.Text)
    subtype = db.Column(db.Text)

    __table_args__ = (db.UniqueConstraint(company, subtype),)

View:

def create():
    instance = Product(**request.json)
    db.session.add(instance)
    try:
        db.session.commit()
    except IntegrityError:
        db.session.rollback()
        return {"detail": "Product object already exists", "status": 406, "title": "Duplicate object"}, 406
    return {"uri": f"/products/{instance.id}"}, 201

Tests:

DEFAULT_DATA = {"company": "Test", "subtype": "Sub"}


def test_create(client):
    response = client.post("/products", json=DEFAULT_DATA)
    assert response.status_code == 201
    instance = Product.query.one()
    assert response.json == {"uri": f"/products/{instance.id}"}


def test_create_duplicate(client):
    response = client.post("/products", json=DEFAULT_DATA)
    assert response.status_code == 201

    instance = Product.query.one()
    assert response.json == {"uri": f"/products/{instance.id}"}
    response = client.post("/products", json=DEFAULT_DATA)
    assert response.status_code == 406
    assert response.json == {"detail": "Product object already exists", "status": 406, "title": "Duplicate object"}

conftest.py:

import flask_migrate
import pytest
from sqlalchemy import event

from project.app import create_connexion_app
from project.models import db


@pytest.fixture(scope="session")
def connexion_app():
    return create_connexion_app("project.settings.TestSettings")


@pytest.fixture(scope="session")
def app(connexion_app):
    app = connexion_app.app
    with app.app_context():
        yield app


@pytest.fixture(scope="session", name="db")
def db_setup(app):
    flask_migrate.upgrade()

    yield db

    flask_migrate.downgrade()
    db.engine.execute("DROP TABLE IF EXISTS alembic_version")


@pytest.fixture(autouse=True)
def session(db):
    with db.engine.connect() as connection:

        @event.listens_for(db.session, "after_transaction_end")
        def restart_savepoint(session, transaction):
            if transaction.nested and not transaction._parent.nested:
                # ensure that state is expired the way
                # session.commit() at the top level normally does
                # (optional step)
                session.expire_all()
                session.begin_nested()

        transaction = connection.begin()

        db.session.begin_nested()

        yield db.session

        db.session.rollback()
        db.session.close()

        if transaction.is_active:
            transaction.rollback()

SQLALCHEMY_COMMIT_ON_TEARDOWN is set to False

The second test is failing with the following output:

    def test_create_duplicate(client):
        response = client.post("/products", json=DEFAULT_DATA)
>       assert response.status_code == 201
E       AssertionError: assert 406 == 201
E        +  where 406 = <<class 'pytest_flask.plugin.JSONResponse'> streamed [406 NOT ACCEPTABLE]>.status_code

Relevant PG log:

LOG:  statement: BEGIN
LOG:  statement: INSERT INTO product (company, subtype) VALUES ('Test', 'Sub') RETURNING product.id
LOG:  statement: COMMIT
LOG:  statement: BEGIN
LOG:  statement: SELECT product.id AS product_id, product.company AS product_company, product.subtype AS product_subtype
    FROM product
    WHERE product.id = 1
LOG:  statement: SELECT product.id AS product_id, product.company AS product_company, product.subtype AS product_subtype
    FROM product
LOG:  statement: ROLLBACK
LOG:  statement: BEGIN
LOG:  statement: INSERT INTO product (company, subtype) VALUES ('Test', 'Sub') RETURNING product.id
ERROR:  duplicate key value violates unique constraint "product_company_subtype_key"
DETAIL:  Key (company, subtype)=(Test, Sub) already exists.
STATEMENT:  INSERT INTO product (company, subtype) VALUES ('Test', 'Sub') RETURNING product.id
LOG:  statement: ROLLBACK

So, the first test commits a row into the DB and it is not rolled back between tests, thus db state is not restored between runs.

Other tests, without involving an explicit rollback work fine. Tried to change SQLALCHEMY_COMMIT_ON_TEARDOWN to True and use flush instead of commit, but in this case the tests after test_create_duplicate are affected.

How to setup a test suite for testing such code, that involves manual commiting / rolling back?

Packages:

• Flask==1.0.2
• Flask-Migrate==2.2.1
• Flask-SQLAlchemy==2.3.2
• SQLAlchemy==1.2.9
• dictalchemy==0.1.2.7
• connexion==1.4.2
• pytest==3.6.2
• pytest-flask==0.10.0

Python version: 3.6.6

RDBMS: PostgreSQL 10.4

Answer 1

python
@pytest.fixture(autouse=True)
def session(db):
    # Begin a top-level transaction
    connection = db.engine.connect()
    transaction = connection.begin()

    # Bind the connection to the session
    db.session.bind = connection

    # Start a nested transaction for the test
    db.session.begin_nested()

    # Listen for the "after_transaction_end" event to reset nested transactions
    @event.listens_for(db.session, "after_transaction_end")
    def restart_savepoint(session, transaction):
        if transaction.nested and not transaction._parent.nested:
            session.begin_nested()

    yield db.session  # Provide the session to the test

    # Roll back the nested transaction and close the session
    db.session.rollback()
    db.session.close()

    # Roll back the top-level transaction
    transaction.rollback()
    connection.close()