CODEWITHSUNDEEP
pythonunit-testingsqlalchemy
gotgenes
gotgenes

Reputation: 40059

Unit tests for Query in SQLAlchemy

How does one go about testing queries in SQLAlchemy? For example suppose we have this models.py

from sqlalchemy import (
        Column,
        Integer,
        String,
)
from sqlalchemy.ext.declarative import declarative_base

Base = declarative_base()

class Panel(Base):
    __tablename__ = 'Panels'

    id = Column(Integer, primary_key=True)
    category = Column(Integer, nullable=False)
    platform = Column(String, nullable=False)
    region = Column(String, nullable=False)

    def __init__(self, category, platform, region):
        self.category = category
        self.platform = platform
        self.region = region


    def __repr__(self):
        return (
            "<Panel('{self.category}', '{self.platform}', "
            "'{self.region}')>".format(self=self)
        )

and this tests.py

import unittest

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from models import Base, Panel


class TestQuery(unittest.TestCase):

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

    def setUp(self):
        Base.metadata.create_all(self.engine)
        self.session.add(Panel(1, 'ion torrent', 'start'))
        self.session.commit()

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

    def test_query_panel(self):
        expected = [Panel(1, 'ion torrent', 'start')]
        result = self.session.query(Panel).all()
        self.assertEqual(result, expected)

When we try running the test, it fails, even though the two Panels look identical.

$ nosetests
F
======================================================================
FAIL: test_query_panel (tests.TestQuery)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/clasher/tmp/tests.py", line 31, in test_query_panel
    self.assertEqual(result, expected)
AssertionError: Lists differ: [<Panel('1', 'ion torrent', 's... != [<Panel('1', 'ion torrent', 's...

First differing element 0:
<Panel('1', 'ion torrent', 'start')>
<Panel('1', 'ion torrent', 'start')>

  [<Panel('1', 'ion torrent', 'start')>, <Panel('2', 'ion torrent', 'end')>]

----------------------------------------------------------------------
Ran 1 test in 0.063s

FAILED (failures=1)

One solution I've found is to make a query for every single instance I expect to find in the query:

class TestQuery(unittest.TestCase):

    ...

    def test_query_panel(self):
        expected = [
            (1, 'ion torrent', 'start'),
            (2, 'ion torrent', 'end')
        ]
        successful = True
        # Check to make sure every expected item is in the query
        try:
            for category, platform, region in expected:
                self.session.query(Panel).filter_by(
                        category=category, platform=platform,
                        region=region).one()
        except (NoResultFound, MultipleResultsFound):
            successful = False
        self.assertTrue(successful)
        # Check to make sure no unexpected items are in the query
        self.assertEqual(self.session.query(Panel).count(),
                         len(expected))

This strikes me as pretty ugly, though, and I'm not even getting to the point where I have a complex filtered query that I'm trying to test. Is there a more elegant solution, or do I always have to manually make a bunch of individual queries?

Upvotes: 43

Views: 52616

Answers (1)

zzzeek
zzzeek

Reputation: 75297

your original test is on the right track, you just have to do one of two things: either make sure that two Panel objects of the same primary key identity compare as True:

import unittest

from sqlalchemy import create_engine
from sqlalchemy.orm import Session

from database.models import Base

class Panel(Base):
    # ...

    def __eq__(self, other):
        return isinstance(other, Panel) and other.id == self.id

or you can organize your test such that you make sure you're checking against the same Panel instance (because here we take advantage of the identity map):

class TestQuery(unittest.TestCase):
    def setUp(self):
        self.engine = create_engine('sqlite:///:memory:')
        self.session = Session(self.engine)
        Base.metadata.create_all(self.engine)
        self.panel = Panel(1, 'ion torrent', 'start')
        self.session.add(self.panel)
        self.session.commit()

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

    def test_query_panel(self):
        expected = [self.panel]
        result = self.session.query(Panel).all()
        self.assertEqual(result, expected)

as far as the engine/session setup/teardown, I'd go for a pattern where you use a single engine for all tests, and assuming your schema is fixed, a single schema for all tests, then you make sure the data you work with is performed within a transaction that can be rolled back. The Session can be made to work this way, such that calling commit() doesn't actually commit the "real" transaction, by wrapping the whole test within an explicit Transaction. The example at https://docs.sqlalchemy.org/en/latest/orm/session_transaction.html#joining-a-session-into-an-external-transaction-such-as-for-test-suites illustrates this usage. Having a ":memory:" engine on every test fixture will take up a lot of memory and not really scale out to other databases besides SQLite.

Upvotes: 40

Related Questions