Database testing with transactions

test_db.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import pytest
import sqlalchemy as sa
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

# In `db.py` or something
Base = declarative_base()
Session = sessionmaker()
DBSession = Session()


class Test(Base):
    __tablename__ = 'test'

    id = sa.Column(sa.Integer, primary_key=True)
    value = sa.Column(sa.Float)


# Possibly in `conftest.py`
@pytest.fixture(scope='session')
def testengine(request):
    engine = sa.create_engine('sqlite:///:memory:')
    Session.configure(bind=engine)
    Base.metadata.create_all(engine)

    def teardown():
        Base.metadata.drop_all(engine)

    request.addfinalizer(teardown)
    return engine


@pytest.fixture
def db_transaction(request, testengine):
    connection = testengine.connect()
    transaction = connection.begin()
    Session.configure(bind=connection)

    def teardown():
        transaction.rollback()
        connection.close()

    request.addfinalizer(teardown)

    return connection


@pytest.fixture
def session(db_transaction):
    return Session()


# Tests
def test_stuff(session, db_transaction):
    t = Test(value=15)
    session.add(t)
    session.commit()
    assert session.query(Test).count() == 1


def test_more_stuff(session, db_transaction):
    c = session.query(Test).count()
    assert c == 0