Django test database runner as a context manager

In my last post I mentioned it might be an idea to wrap up the Django test database setup / teardown in a context manager for use with Python’s with statement. Here’s my first stab, which seems to work.

from contextlib import contextmanager


@contextmanager
def test_db_connection():
    """A context manager for Django's test runner.

    For Python 2.5 you will need
        from __future__ import with_statement
    """

    from django.conf import settings
    from django.test.utils import setup_test_environment, teardown_test_environment
    from django.db import connection

    setup_test_environment()

    settings.DEBUG = False    
    verbosity = 0
    interactive = False

    old_name = settings.DATABASE_NAME
    connection.creation.create_test_db(verbosity, autoclobber=not interactive)

    yield connection

    connection.creation.destroy_test_db(old_name, verbosity)
    teardown_test_environment()

All of this requires Python 2.5 or later.

So with that snippet you could write a test something like so:

import unittest


class MyTestCase(unittest.TestCase):
    def test_myModelTest(self):
        with test_db_connection():
            from myproject.myapp.models import MyModel

            obj = MyModel()
            obj.save()
            self.assert_(obj.pk)

… and just as with Django’s manage.py test command the objects would be created within the test database then destroyed when the with test_db_connection() block is finished.

Everything’s going to be hunky dory.

Leave a Reply

Your email address will not be published. Required fields are marked *