Test case now drops the mongo collection

cms/djangoapps/contentstore/tests/utils.py

@@ -1,3 +1,9 @@
+'''
+Utilities for contentstore tests
+'''
+
+#pylint: disable=W0603
+
 import json
 import copy
 from uuid import uuid4
@@ -17,37 +23,90 @@ class ModuleStoreTestCase(TestCase):
     collection with templates before running the TestCase
     and drops it they are finished. """
 
-    def _pre_setup(self):
-        super(ModuleStoreTestCase, self)._pre_setup()
+    @staticmethod
+    def flush_mongo_except_templates():
+        '''
+        Delete everything in the module store except templates
+        '''
+        modulestore = xmodule.modulestore.django.modulestore()
+
+        # This query means: every item in the collection
+        # that is not a template
+        query = { "_id.course": { "$ne": "templates" }}
+
+        # Remove everything except templates
+        modulestore.collection.remove(query)
+
+    @staticmethod
+    def load_templates_if_necessary():
+        '''
+        Load templates into the modulestore only if they do not already exist.
+        We need the templates, because they are copied to create
+        XModules such as sections and problems
+        '''
+        modulestore = xmodule.modulestore.django.modulestore()
+
+        # Count the number of templates
+        query = { "_id.course": "templates"}
+        num_templates = modulestore.collection.find(query).count()
+
+        if num_templates < 1:
+            update_templates()
+
+    @classmethod
+    def setUpClass(cls):
+        '''
+        Flush the mongo store and set up templates
+        '''
 
         # Use a uuid to differentiate
         # the mongo collections on jenkins.
-        self.orig_MODULESTORE = copy.deepcopy(settings.MODULESTORE)
-        self.test_MODULESTORE = self.orig_MODULESTORE
-        self.test_MODULESTORE['default']['OPTIONS']['collection'] = 'modulestore_%s' % uuid4().hex
-        self.test_MODULESTORE['direct']['OPTIONS']['collection'] = 'modulestore_%s' % uuid4().hex
-        settings.MODULESTORE = self.test_MODULESTORE
-
-        # Flush and initialize the module store
-        # It needs the templates because it creates new records
-        # by cloning from the template.
-        # Note that if your test module gets in some weird state
-        # (though it shouldn't), do this manually
-        # from the bash shell to drop it:
-        # $ mongo test_xmodule --eval "db.dropDatabase()"
+        cls.orig_modulestore = copy.deepcopy(settings.MODULESTORE)
+        test_modulestore = cls.orig_modulestore
+        test_modulestore['default']['OPTIONS']['collection'] = 'modulestore_%s' % uuid4().hex
+        test_modulestore['direct']['OPTIONS']['collection'] = 'modulestore_%s' % uuid4().hex
         xmodule.modulestore.django._MODULESTORES = {}
-        update_templates()
+
+        settings.MODULESTORE = test_modulestore
+
+        TestCase.setUpClass()
+
+    @classmethod
+    def tearDownClass(cls):
+        '''
+        Revert to the old modulestore settings
+        '''
+
+        # Clean up by dropping the collection
+        modulestore = xmodule.modulestore.django.modulestore()
+        modulestore.collection.drop()
+
+        # Restore the original modulestore settings
+        settings.MODULESTORE = cls.orig_modulestore
+
+    def _pre_setup(self):
+        '''
+        Remove everything but the templates before each test
+        '''
+
+        # Flush anything that is not a template
+        ModuleStoreTestCase.flush_mongo_except_templates()
+
+        # Check that we have templates loaded; if not, load them
+        ModuleStoreTestCase.load_templates_if_necessary()
+
+        # Call superclass implementation
+        TestCase._pre_setup(self)
 
     def _post_teardown(self):
-        # Make sure you flush out the modulestore.
-        # Drop the collection at the end of the test,
-        # otherwise there will be lingering collections leftover
-        # from executing the tests.
-        xmodule.modulestore.django._MODULESTORES = {}
-        xmodule.modulestore.django.modulestore().collection.drop()
-        settings.MODULESTORE = self.orig_MODULESTORE
+        '''
+        Flush everything we created except the templates
+        '''
+        # Flush anything that is not a template
+        ModuleStoreTestCase.flush_mongo_except_templates()
 
-        super(ModuleStoreTestCase, self)._post_teardown()
+        # Call superclass implementation
+        TestCase._post_teardown(self)
 
 
 def parse_json(response):