Skip to content

Instantly share code, notes, and snippets.

@andgineer
Created June 9, 2020 15:38
Show Gist options
  • Star 10 You must be signed in to star a gist
  • Fork 2 You must be signed in to fork a gist
  • Save andgineer/a13f2aeb7b46c6269a9ec7d09edad91c to your computer and use it in GitHub Desktop.
Save andgineer/a13f2aeb7b46c6269a9ec7d09edad91c to your computer and use it in GitHub Desktop.
pytest Fixture to rollback all DB changes after test
from sqlalchemy import event
from sqlalchemy.orm import sessionmaker
from app.db.session import engine
import pytest
import app.tests.config
@pytest.fixture(
scope='function',
autouse=True # New test DB session for each test todo we need it only for tests with Client fixture
)
def db():
"""
SQLAlchemy session started with SAVEPOINT
After test rollback to this SAVEPOINT
"""
connection = engine(
app.tests.config.get_test_config()
).connect()
# begin a non-ORM transaction
trans = connection.begin()
session = sessionmaker()(bind=connection)
session.begin_nested() # SAVEPOINT
app.tests.config.session = session # Inject session to the server code under test
@event.listens_for(app.tests.config.session, "after_transaction_end")
def restart_savepoint(session, transaction):
"""
Each time that SAVEPOINT ends, reopen it
"""
if transaction.nested and not transaction._parent.nested:
session.begin_nested()
yield session
session.close()
trans.rollback() # roll back to the SAVEPOINT
connection.close()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment