Custom sqlite database for unit tests for code using peewee ORM

Question

I am trying to implement a many-to-many scenario using peewee python ORM and I'd like some unit tests. Peewee tutorial is great but it assumes that database is defined at module level then all models are using it. My situation is different: I don't have a source code file (a module from python's point of view) with tests which I run explicitly, I am using nose which collects tests from that file and runs them.

How do I use a custom database only for models instantiated in tests (which are being run by nose)? My goal is to use an in-memory database for tests only, to speedup the testing process.

Answer 1

I just pushed a commit today that makes this easier.

The fix is in the form of a context manager which allows you to override the database of a model:

from unittest import TestCase
from playhouse.test_utils import test_database
from peewee import *

from my_app.models import User, Tweet

test_db = SqliteDatabase(':memory:')

class TestUsersTweets(TestCase):
    def create_test_data(self):
        # ... create a bunch of users and tweets
        for i in range(10):
            User.create(username='user-%d' % i)

    def test_timeline(self):
        with test_database(test_db, (User, Tweet)):
            # This data will be created in `test_db`
            self.create_test_data()

            # Perform assertions on test data inside ctx manager.
            self.assertEqual(Tweet.timeline('user-0') [...])

        # once we exit the context manager, we're back to using the normal database

See the documentation and have a look at the example testcases:

• Context manager
• Testcases showing how to use

Answer 2

To not include context manager in every test case, overwrite run method.

# imports and db declaration

class TestUsersTweets(TestCase):
    def run(self, result=None):
        with test_database(test_db, (User, Tweet)):
            super(TestUsersTweets, self).run(result)

    def test_timeline(self):
        self.create_test_data()
        self.assertEqual(Tweet.timeline('user-0') [...])

Answer 3

I took the great answers from @coleifer and @avalanchy and took them one step further.

In order to avoid overriding the run method on every TestCase subclass, you can use a base class... and I also like the idea of not having to write down every model class I work with, so I came up with this

import unittest
import inspect
import sys
import peewee
from abc import ABCMeta
from playhouse.test_utils import test_database
from business_logic.models import *

test_db = peewee.SqliteDatabase(':memory:')


class TestCaseWithPeewee(unittest.TestCase):
    """
    This abstract class is used to "inject" the test database so that the tests don't use the real sqlite db
    """

    __metaclass__ = ABCMeta

    def run(self, result=None):
        model_classes = [m[1] for m in inspect.getmembers(sys.modules['business_logic.models'], inspect.isclass) if
                         issubclass(m[1], peewee.Model) and m[1] != peewee.Model]
        with test_database(test_db, model_classes):
            super(TestCaseWithPeewee, self).run(result)

so, now I can just inherit from TestCaseWithPeewee and don't have to worry about anything else other than the test

Answer 4

Apparently, there's a new approach for the scenario described, where you can bind the models in the setUp() method of your test case:

Example from the official docs:

# tests.py
import unittest
from my_app.models import EventLog, Relationship, Tweet, User

MODELS = [User, Tweet, EventLog, Relationship]

# use an in-memory SQLite for tests.
test_db = SqliteDatabase(':memory:')

class BaseTestCase(unittest.TestCase):
    def setUp(self):
        # Bind model classes to test db. Since we have a complete list of
        # all models, we do not need to recursively bind dependencies.
        test_db.bind(MODELS, bind_refs=False, bind_backrefs=False)

        test_db.connect()
        test_db.create_tables(MODELS)

    def tearDown(self):
        # Not strictly necessary since SQLite in-memory databases only live
        # for the duration of the connection, and in the next step we close
        # the connection...but a good practice all the same.
        test_db.drop_tables(MODELS)

        # Close connection to db.
        test_db.close()

        # If we wanted, we could re-bind the models to their original
        # database here. But for tests this is probably not necessary.

Answer 5

When using test_database I encountered problems with test_db not being initialized:

nose.proxy.Exception: Error, database not properly initialized before opening connection -------------------- >> begin captured logging << -------------------- peewee: DEBUG: ('SELECT "t1"."id", "t1"."name", "t1"."count" FROM "counter" AS t1', []) --------------------- >> end captured logging << ---------------------

I eventually fixed this by passing create_tables=True like so:

def test_timeline(self): with test_database(test_db, (User, Tweet), create_tables=True): # This data will be created in `test_db` self.create_test_data()

According to the docs create_tables should default to True but it seems that isn't the case in the latest release of peewee.

Answer 6

For anyone who's using pytest, here's how I did it:

conftest.py

MODELS = [User, Tweet]  # Also add get_through_model() for ManyToMany fields

test_db = SqliteDatabase(':memory:')
test_db.bind(MODELS, bind_refs=False, bind_backrefs=False)
test_db.connect()
test_db.create_tables(MODELS)

@pytest.fixture(autouse=True)
def in_mem_db(mocker):
    mocked_db = mocker.patch("database.db", autospec=True)  # "database.db" is where your app's code imports db from
    mocked_db.return_value = test_db
    return mocked_db

And voila, all your tests run with an in-memory sqlite database.