Mongoengine Test Case Parent for Django

test_utilities.py

from mongoengine.connection import connect, disconnect

from django.test import TestCase
from django.conf import settings


class MongoTestCase(TestCase):
    """
    TestCase class that clear the collection between the tests
    """
    _connections = {}

    def _pre_setup(self):
        disconnect()
        MONGO_DATABASES = settings.MONGO_DATABASES
        for db_name in MONGO_DATABASES:
            if db_name.find("test") > -1:
                continue
            test_db_name = "test__%s" % db_name
            self._connections[db_name] = connect(test_db_name,
                                                 alias=db_name,
                                                 host=MONGO_DATABASES[db_name]["HOST"],
                                                 port=MONGO_DATABASES[db_name].get("PORT", 27017))
        super(MongoTestCase, self)._pre_setup()

    def _post_teardown(self):
        """
        Loop over each connection. The key name is also the name of a database we want to clear,
        so get that database at self._connections[db_name][db_name] and drop all the non-system
        collections
        """
        for db_name in self._connections:
            db = self._connections[db_name][db_name]
            for collection in [name for name in
                               db.collection_names()
                               if name.find("system.") == -1]:
                db[collection].drop()
            disconnect(db_name)
        super(MongoTestCase, self)._post_teardown()