Custom TestRunner and APITestCase to work with django, mongodb and django rest framework

django.mongo.testsutils.py

from django.test.runner import DiscoverRunner
from django.conf import settings
from mongoengine import connect, disconnect, get_connection
from pymongo import MongoClient
from rest_framework.test import APITestCase

class TestRunner(DiscoverRunner):
    """
    Custom test runner for mongoengine
    """
    def setup_databases(self, **kwargs):
        db_name = settings.TEST_DBNAME
        disconnect()
        connect(db_name)
        print('Connecting to test database: ' + db_name)
        return db_name

    def teardown_databases(self, db_name, **kwargs):
        if not self.keepdb:
            conn = MongoClient()
            conn.drop_database(db_name)
            print('Dropping test database: ' + db_name)


class MongoTestCase(APITestCase):
    """
    Custom APITestCase that clear the database between the tests
    """
    db_name = settings.TEST_DBNAME
    
    def _pre_setup(self):
        disconnect()
        connect(self.db_name)
        super(MongoTestCase, self)._pre_setup()

    def _post_teardown(self):
        connection = get_connection()
        connection.drop_database(self.db_name)
        disconnect()
        super(MongoTestCase, self)._post_teardown()