Search This Blog

05 August 2014

Generate unit tests for sqlalchemy mappings

When playing with sqlalchemy mappings, it seems a good idea do unit test them, just to make sure they compile properly. So, given this simple mappings (from sqla tutorial)
__author__ = 'nlaurance'
from sqlalchemy import Column, DateTime, String, Integer, ForeignKey, func
from sqlalchemy.orm import relationship, backref
from sqlalchemy.ext.declarative import declarative_base


Base = declarative_base()


class Department(Base):
    __tablename__ = 'department'
    id = Column(Integer, primary_key=True)
    name = Column(String, nullable=False)


class Employee(Base):
    __tablename__ = 'employee'
    id = Column(Integer, primary_key=True)
    name = Column(String)
    hired_on = Column(DateTime, default=func.now())
    department_id = Column(Integer, ForeignKey('department.id'))
    department = relationship(
        Department,
        backref=backref('employees',
                         uselist=True,
                         cascade='delete,all'))

it would make sense to, at least make sure we can create one and retrieve its instance Let's create a test base, that our tests will inherit from, and factor out sensible defaults that could be used in integration tests later on.
__author__ = 'nlaurance'
import inspect
from unittest import TestCase
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqla_classes import Base
from sqla_classes import Department, Employee


DEFAULTS = {
    Department: {
        "name": "Lorem",
    },
    Employee: {
        "name": "Lorem",
    },
}


def create_some(session, klass, values=None):
    default = DEFAULTS.get(klass, {})
    values = values or [default]
    for value in values:
        new_default = default.copy()
        new_default.update(value)
        session.add(klass(**new_default))


class SqlAMappingsBase(object):

    engine = create_engine('sqlite:///:memory:')
    session = sessionmaker(bind=engine)()

    class_under_test = None
    default_values = {}

    def setUp(self):
        Base.metadata.create_all(self.engine)
        create_some(self.session, self.class_under_test)
        self.session.commit()

    def tearDown(self):
        Base.metadata.drop_all(self.engine)

    def test_get(self):
        self.session.query(self.class_under_test).one()
now if we want to unit test Department and Employee classes we would write something like :
class DepartmentTest(SqlAMappingsBase, TestCase):
    class_under_test = Department
    default_values = DEFAULTS.get(Department, {})

class EmployeeTest(SqlAMappingsBase, TestCase):
    class_under_test = Employee
    default_values = DEFAULTS.get(Employee, {})
so far, so good. but ... what if the projects requires to create lots of mappings, we would end up with lots of look alike code, mostly copied/pasted from one of the examples. Why not create the test classes on the fly ? some code like ..
def autotest_mappings(module_under_test, test_module):
    for name, klass in inspect.getmembers(module_under_test, inspect.isclass):
        if Base in inspect.getmro(klass) and klass is not Base:
            test_name = ''.join((name, 'Test'))
            test_klass = type(test_name, (SqlAMappingsBase, TestCase),
                              {"class_under_test": klass,
                               "default_values": DEFAULTS.get(klass, {})})
            setattr(test_module, test_name, test_klass)
this function takes as arguments the module we want to test, and the test module in which to insert the tests. It then looks for all classes that inherit from Base, create a new class and inject it into the test module. Our test code would then look like :
from our_project import models
from sqla_utils import autotest_mappings

autotest_mappings(models, sys.modules[__name__])
and that's it ! three lines and no matter how many classes you declare in your models module, each would be picked by the inspect sniffer and would have a nice unit test.