[Python] How to test rollbacks in SQLAlchemy

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

Continue reading...