Context managers are sexy for testing.

testcase.py

# -*- coding: utf-8 -*-

from django.test import TestCase


class _ObjectContext(object):
    u"""
    Context manager returned by assertObjectCreated/assertObjectDeleted.
    """
    def __init__(self, enter_assertion, exit_assertion, model_class, **kwargs):
        self.enter_assertion = enter_assertion
        self.exit_assertion = exit_assertion
        self.model_class = model_class
        self.kwargs = kwargs

    def __enter__(self):
        self.enter_assertion(self.model_class, **self.kwargs)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.exit_assertion(self.model_class, **self.kwargs)
        return True


class ObjectTestCaseMixin(object):
    u"""
    Adds ORM-related assertions for testing object creation and deletion.
    """
    def assertObjectExists(self, model_class, **kwargs):
        u"""
        Makes sure that an object exists.
        """
        try:
            obj = model_class._default_manager.get(**kwargs)
            self.assertIsNotNone(obj)
        except model_class.DoesNotExist:
            self.fail(u"No %s found matching the criteria." % model_class)

    def assertObjectDoesNotExist(self, model_class, **kwargs):
        u"""
        Asserts that an object does not exist.
        """
        try:
            instance = model_class._default_manager.get(**kwargs)
            self.fail(u"A %s was found matching the criteria. (%s)" % (
                model_class,
                instance,
            ))
        except model_class.DoesNotExist:
            pass

    def assertObjectCreated(self, model_class, **kwargs):
        u"""
        A context manager to help checking if the object is created.

        Example::

            with self.assertObjectCreated(Article, slug='lorem-ipsum'):
                Article.objects.create(slug='lorem-ipsum')
        """
        return _ObjectContext(self.assertObjectDoesNotExist,
                              self.assertObjectExists,
                              model_class,
                              **kwargs)

    def assertObjectDeleted(self, model_class, **kwargs):
        u"""
        A context manager to help checking if the object is deleted.

        Example::

            with self.assertObjectDeleted(Article, slug='lorem-ipsum'):
                article = get_object_or_404(Article, slug='lorem-ipsum')
                article.delete()
        """
        return _ObjectContext(self.assertObjectExists,
                              self.assertObjectDoesNotExist,
                              model_class,
                              **kwargs)


class BaseTestCase(TestCase, ObjectTestCaseMixin):
    pass