Merge pull request #6920 from edx/jeskew/eager_loading_xblocks_to_depth

Improve performance of import/export and course traversal with Split modulestore.

common/lib/xmodule/xmodule/modulestore/perf_tests/test_asset_import_export.py

@@ -105,14 +105,8 @@ class CrossStoreXMLRoundtrip(unittest.TestCase):
                 make_asset_xml(num_assets, ASSET_XML_PATH)
                 validate_xml(ASSET_XSD_PATH, ASSET_XML_PATH)
 
-            # Construct the contentstore for storing the first import
-            with MongoContentstoreBuilder().build() as source_content:
-                # Construct the modulestore for storing the first import (using the previously created contentstore)
-                with source_ms.build(source_content) as source_store:
-                    # Construct the contentstore for storing the second import
-                    with MongoContentstoreBuilder().build() as dest_content:
-                        # Construct the modulestore for storing the second import (using the second contentstore)
-                        with dest_ms.build(dest_content) as dest_store:
+            with source_ms.build() as (source_content, source_store):
+                with dest_ms.build() as (dest_content, dest_store):
                     source_course_key = source_store.make_course_key('a', 'course', 'course')
                     dest_course_key = dest_store.make_course_key('a', 'course', 'course')
 
@@ -193,10 +187,7 @@ class FindAssetTest(unittest.TestCase):
                 make_asset_xml(num_assets, ASSET_XML_PATH)
                 validate_xml(ASSET_XSD_PATH, ASSET_XML_PATH)
 
-            # Construct the contentstore for storing the first import
-            with MongoContentstoreBuilder().build() as source_content:
-                # Construct the modulestore for storing the first import (using the previously created contentstore)
-                with source_ms.build(source_content) as source_store:
+            with source_ms.build() as (source_content, source_store):
                 source_course_key = source_store.make_course_key('a', 'course', 'course')
                 asset_key = source_course_key.make_asset_key(
                     AssetMetadata.GENERAL_ASSET_TYPE, 'silly_cat_picture.gif'
@@ -265,10 +256,7 @@ class TestModulestoreAssetSize(unittest.TestCase):
         make_asset_xml(num_assets, ASSET_XML_PATH)
         validate_xml(ASSET_XSD_PATH, ASSET_XML_PATH)
 
-        # Construct the contentstore for storing the first import
-        with MongoContentstoreBuilder().build() as source_content:
-            # Construct the modulestore for storing the first import (using the previously created contentstore)
-            with source_ms.build(source_content) as source_store:
+        with source_ms.build() as (source_content, source_store):
             source_course_key = source_store.make_course_key('a', 'course', 'course')
 
             import_from_xml(

common/lib/xmodule/xmodule/modulestore/split_mongo/split.py

@@ -402,13 +402,18 @@ class SplitBulkWriteMixin(BulkOperationsMixin):
 
         bulk_write_record = self._get_bulk_ops_record(course_key)
         if bulk_write_record.active:
+            # Only query for the definitions that aren't already cached.
             for definition in bulk_write_record.definitions.values():
                 definition_id = definition.get('_id')
                 if definition_id in ids:
                     ids.remove(definition_id)
                     definitions.append(definition)
 
-        definitions.extend(self.db_connection.get_definitions(list(ids)))
+        # Query the db for the definitions.
+        defs_from_db = self.db_connection.get_definitions(list(ids))
+        # Add the retrieved definitions to the cache.
+        bulk_write_record.definitions.update({d.get('_id'): d for d in defs_from_db})
+        definitions.extend(defs_from_db)
         return definitions
 
     def update_definition(self, course_key, definition):
@@ -683,23 +688,29 @@ class SplitMongoModuleStore(SplitBulkWriteMixin, ModuleStoreWriteBase):
                     new_module_data
                 )
 
-            if not lazy:
-                # Load all descendants by id
+            # This code supports lazy loading, where the descendent definitions aren't loaded
+            # until they're actually needed.
+            # However, assume that depth == 0 means no depth is specified and depth != 0 means
+            # a depth *is* specified. If a non-zero depth is specified, force non-lazy definition
+            # loading in order to populate the definition cache for later access.
+            load_definitions_now = depth != 0 or not lazy
+            if load_definitions_now:
+                # Non-lazy loading: Load all descendants by id.
                 descendent_definitions = self.get_definitions(
                     course_key,
                     [
-                        block['definition']
+                        block.definition
                         for block in new_module_data.itervalues()
                     ]
                 )
-                # turn into a map
+                # Turn definitions into a map.
                 definitions = {definition['_id']: definition
                                for definition in descendent_definitions}
 
                 for block in new_module_data.itervalues():
                     if block.definition in definitions:
                         definition = definitions[block.definition]
-                        # convert_fields was being done here, but it gets done later in the runtime's xblock_from_json
+                        # convert_fields gets done later in the runtime's xblock_from_json
                         block.fields.update(definition.get('fields'))
                         block.definition_loaded = True
 

common/lib/xmodule/xmodule/modulestore/tests/test_asides.py

@@ -34,7 +34,7 @@ class TestAsidesXmlStore(TestCase):
         """
         Check that the xml modulestore read in all the asides with their values
         """
-        with XmlModulestoreBuilder().build(course_ids=['edX/aside_test/2012_Fall']) as store:
+        with XmlModulestoreBuilder().build(course_ids=['edX/aside_test/2012_Fall']) as (__, store):
             def check_block(block):
                 """
                 Check whether block has the expected aside w/ its fields and then recurse to the block's children

common/lib/xmodule/xmodule/modulestore/tests/test_assetstore.py

@@ -161,8 +161,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Save the metadata in each store and retrieve it singularly, as all assets, and after deleting all.
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             asset_filename = 'burnside.jpg'
             new_asset_loc = course.id.make_asset_key('asset', asset_filename)
@@ -180,8 +179,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Delete non-existent and existent metadata
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             new_asset_loc = course.id.make_asset_key('asset', 'burnside.jpg')
             # Attempt to delete an asset that doesn't exist.
@@ -198,8 +196,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Find a non-existent asset in an existing course.
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             new_asset_loc = course.id.make_asset_key('asset', 'burnside.jpg')
             # Find existing asset metadata.
@@ -211,8 +208,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Get all assets in an existing course when no assets exist.
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             # Find existing asset metadata.
             asset_md = store.get_all_asset_metadata(course.id, 'asset')
@@ -223,8 +219,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Find asset metadata from a non-existent course.
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             fake_course_id = CourseKey.from_string("{}nothere/{}nothere/{}nothere".format(
                 course.id.org, course.id.course, course.id.run
@@ -241,8 +236,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Add an asset's metadata, then add it again.
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             new_asset_loc = course.id.make_asset_key('asset', 'burnside.jpg')
             new_asset_md = self._make_asset_metadata(new_asset_loc)
@@ -259,8 +253,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Test saving assets with other asset types.
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             new_asset_loc = course.id.make_asset_key('vrml', 'pyramid.vrml')
             new_asset_md = self._make_asset_metadata(new_asset_loc)
@@ -274,8 +267,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Test saving assets using an asset type of 'course_id'.
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             new_asset_loc = course.id.make_asset_key('course_id', 'just_to_see_if_it_still_works.jpg')
             new_asset_md = self._make_asset_metadata(new_asset_loc)
@@ -291,8 +283,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Save multiple metadata in each store and retrieve it singularly, as all assets, and after deleting all.
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             new_asset_loc = course.id.make_asset_key('asset', 'burnside.jpg')
             new_asset_md = self._make_asset_metadata(new_asset_loc)
@@ -340,71 +331,67 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Save setting each attr one at a time
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             new_asset_loc = course.id.make_asset_key('asset', 'burnside.jpg')
             new_asset_md = self._make_asset_metadata(new_asset_loc)
             store.save_asset_metadata(new_asset_md, ModuleStoreEnum.UserID.test)
-                for attr, value in self.ALLOWED_ATTRS:
-                    # Set the course asset's attr.
-                    store.set_asset_metadata_attr(new_asset_loc, attr, value, ModuleStoreEnum.UserID.test)
-                    # Find the same course asset and check its changed attr.
+            for attribute, value in self.ALLOWED_ATTRS:
+                # Set the course asset's attribute.
+                store.set_asset_metadata_attr(new_asset_loc, attribute, value, ModuleStoreEnum.UserID.test)
+                # Find the same course asset and check its changed attribute.
                 updated_asset_md = store.find_asset_metadata(new_asset_loc)
                 self.assertIsNotNone(updated_asset_md)
-                    self.assertIsNotNone(getattr(updated_asset_md, attr, None))
-                    self.assertEquals(getattr(updated_asset_md, attr, None), value)
+                self.assertIsNotNone(getattr(updated_asset_md, attribute, None))
+                self.assertEquals(getattr(updated_asset_md, attribute, None), value)
 
     @ddt.data(*MODULESTORE_SETUPS)
     def test_set_disallowed_attrs(self, storebuilder):
         """
         setting disallowed attrs should fail
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             new_asset_loc = course.id.make_asset_key('asset', 'burnside.jpg')
             new_asset_md = self._make_asset_metadata(new_asset_loc)
             store.save_asset_metadata(new_asset_md, ModuleStoreEnum.UserID.test)
-                for attr, value in self.DISALLOWED_ATTRS:
-                    original_attr_val = getattr(new_asset_md, attr)
-                    # Set the course asset's attr.
-                    store.set_asset_metadata_attr(new_asset_loc, attr, value, ModuleStoreEnum.UserID.test)
-                    # Find the same course and check its changed attr.
+            for attribute, value in self.DISALLOWED_ATTRS:
+                original_attr_val = getattr(new_asset_md, attribute)
+                # Set the course asset's attribute.
+                store.set_asset_metadata_attr(new_asset_loc, attribute, value, ModuleStoreEnum.UserID.test)
+                # Find the same course and check its changed attribute.
                 updated_asset_md = store.find_asset_metadata(new_asset_loc)
                 self.assertIsNotNone(updated_asset_md)
-                    self.assertIsNotNone(getattr(updated_asset_md, attr, None))
-                    # Make sure that the attr is unchanged from its original value.
-                    self.assertEquals(getattr(updated_asset_md, attr, None), original_attr_val)
+                self.assertIsNotNone(getattr(updated_asset_md, attribute, None))
+                # Make sure that the attribute is unchanged from its original value.
+                self.assertEquals(getattr(updated_asset_md, attribute, None), original_attr_val)
 
     @ddt.data(*MODULESTORE_SETUPS)
     def test_set_unknown_attrs(self, storebuilder):
         """
         setting unknown attrs should fail
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             new_asset_loc = course.id.make_asset_key('asset', 'burnside.jpg')
             new_asset_md = self._make_asset_metadata(new_asset_loc)
             store.save_asset_metadata(new_asset_md, ModuleStoreEnum.UserID.test)
-                for attr, value in self.UNKNOWN_ATTRS:
-                    # Set the course asset's attr.
-                    store.set_asset_metadata_attr(new_asset_loc, attr, value, ModuleStoreEnum.UserID.test)
-                    # Find the same course and check its changed attr.
+            for attribute, value in self.UNKNOWN_ATTRS:
+                # Set the course asset's attribute.
+                store.set_asset_metadata_attr(new_asset_loc, attribute, value, ModuleStoreEnum.UserID.test)
+                # Find the same course and check its changed attribute.
                 updated_asset_md = store.find_asset_metadata(new_asset_loc)
                 self.assertIsNotNone(updated_asset_md)
                 # Make sure the unknown field was *not* added.
                 with self.assertRaises(AttributeError):
-                        self.assertEquals(getattr(updated_asset_md, attr), value)
+                    self.assertEquals(getattr(updated_asset_md, attribute), value)
 
     @ddt.data(*MODULESTORE_SETUPS)
     def test_save_one_different_asset(self, storebuilder):
         """
         saving and deleting things which are not 'asset'
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             asset_key = course.id.make_asset_key('different', 'burn.jpg')
             new_asset_thumbnail = self._make_asset_thumbnail_metadata(
@@ -420,8 +407,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         finding things which are of type other than 'asset'
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             asset_key = course.id.make_asset_key('different', 'burn.jpg')
             new_asset_thumbnail = self._make_asset_thumbnail_metadata(
@@ -447,8 +433,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         getting all things which are of type other than 'asset'
         """
         # pylint: disable=bad-continuation
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
 
             # Save 'em.
@@ -485,8 +470,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         Save a list of asset metadata all at once.
         """
         # pylint: disable=bad-continuation
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
 
             # Make a list of AssetMetadata objects.
@@ -526,8 +510,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         Save a list of asset metadata all at once - but with one asset's metadata from a different course.
         """
         # pylint: disable=bad-continuation
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course1 = CourseFactory.create(modulestore=store)
             course2 = CourseFactory.create(modulestore=store)
 
@@ -569,8 +552,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         deleting all assets of a given but not 'asset' type
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course = CourseFactory.create(modulestore=store)
             asset_key = course.id.make_asset_key('different', 'burn_thumb.jpg')
             new_asset_thumbnail = self._make_asset_thumbnail_metadata(
@@ -585,8 +567,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Save multiple metadata in each store and retrieve it singularly, as all assets, and after deleting all.
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course1 = CourseFactory.create(modulestore=store)
             course2 = CourseFactory.create(modulestore=store)
             self.setup_assets(course1.id, course2.id, store)
@@ -659,7 +640,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Test coverage which shows that for now xml read operations are not implemented
         """
-        with storebuilder.build(None) as store:
+        with storebuilder.build(contentstore=None) as (__, store):
             course_key = store.make_course_key("org", "course", "run")
             asset_key = course_key.make_asset_key('asset', 'foo.jpg')
             self.assertEquals(store.find_asset_metadata(asset_key), None)
@@ -670,8 +651,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Create a course with assets, copy them all to another course in the same modulestore, and check on it.
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course1 = CourseFactory.create(modulestore=store)
             course2 = CourseFactory.create(modulestore=store)
             self.setup_assets(course1.id, None, store)
@@ -691,8 +671,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         """
         Create a course with *no* assets, and try copy them all to another course in the same modulestore.
         """
-        with MongoContentstoreBuilder().build() as contentstore:
-            with storebuilder.build(contentstore) as store:
+        with storebuilder.build() as (__, store):
             course1 = CourseFactory.create(modulestore=store)
             course2 = CourseFactory.create(modulestore=store)
             store.copy_all_asset_metadata(course1.id, course2.id, ModuleStoreEnum.UserID.test * 101)
@@ -713,8 +692,7 @@ class TestMongoAssetMetadataStorage(unittest.TestCase):
         Create a course with assets, copy them all to another course in a different modulestore, and check on it.
         """
         mixed_builder = MIXED_MODULESTORE_BOTH_SETUP
-        with MongoContentstoreBuilder().build() as contentstore:
-            with mixed_builder.build(contentstore) as mixed_store:
+        with mixed_builder.build() as (__, mixed_store):
             with mixed_store.default_store(from_store):
                 course1 = CourseFactory.create(modulestore=mixed_store)
             with mixed_store.default_store(to_store):

common/lib/xmodule/xmodule/modulestore/tests/test_cross_modulestore_import_export.py

@@ -76,12 +76,66 @@ class MemoryCache(object):
         self._data[key] = value
 
 
-class MongoModulestoreBuilder(object):
+class MongoContentstoreBuilder(object):
+    """
+    A builder class for a MongoContentStore.
+    """
+    @contextmanager
+    def build(self):
+        """
+        A contextmanager that returns a MongoContentStore, and deletes its contents
+        when the context closes.
+        """
+        contentstore = MongoContentStore(
+            db='contentstore{}'.format(random.randint(0, 10000)),
+            collection='content',
+            **COMMON_DOCSTORE_CONFIG
+        )
+        contentstore.ensure_indexes()
+
+        try:
+            yield contentstore
+        finally:
+            # Delete the created database
+            contentstore._drop_database()  # pylint: disable=protected-access
+
+    def __repr__(self):
+        return 'MongoContentstoreBuilder()'
+
+
+class StoreBuilderBase(object):
+    """
+    Base class for all modulestore builders.
+    """
+    @contextmanager
+    def build(self, **kwargs):
+        """
+        Build the modulstore, optionally building the contentstore as well.
+        """
+        contentstore = kwargs.pop('contentstore', None)
+        if not contentstore:
+            with self.build_without_contentstore() as (contentstore, modulestore):
+                yield contentstore, modulestore
+        else:
+            with self.build_with_contentstore(contentstore) as modulestore:
+                yield modulestore
+
+    @contextmanager
+    def build_without_contentstore(self):
+        """
+        Build both the contentstore and the modulestore.
+        """
+        with MongoContentstoreBuilder().build() as contentstore:
+            with self.build_with_contentstore(contentstore) as modulestore:
+                yield contentstore, modulestore
+
+
+class MongoModulestoreBuilder(StoreBuilderBase):
     """
     A builder class for a DraftModuleStore.
     """
     @contextmanager
-    def build(self, contentstore):
+    def build_with_contentstore(self, contentstore):
         """
         A contextmanager that returns an isolated mongo modulestore, and then deletes
         all of its data at the end of the context.
@@ -125,12 +179,12 @@ class MongoModulestoreBuilder(object):
         return 'MongoModulestoreBuilder()'
 
 
-class VersioningModulestoreBuilder(object):
+class VersioningModulestoreBuilder(StoreBuilderBase):
     """
     A builder class for a VersioningModuleStore.
     """
     @contextmanager
-    def build(self, contentstore):
+    def build_with_contentstore(self, contentstore):
         """
         A contextmanager that returns an isolated versioning modulestore, and then deletes
         all of its data at the end of the context.
@@ -170,13 +224,13 @@ class VersioningModulestoreBuilder(object):
         return 'SplitModulestoreBuilder()'
 
 
-class XmlModulestoreBuilder(object):
+class XmlModulestoreBuilder(StoreBuilderBase):
     """
     A builder class for a XMLModuleStore.
     """
     # pylint: disable=unused-argument
     @contextmanager
-    def build(self, contentstore=None, course_ids=None):
+    def build_with_contentstore(self, contentstore=None, course_ids=None):
         """
         A contextmanager that returns an isolated xml modulestore
 
@@ -194,7 +248,7 @@ class XmlModulestoreBuilder(object):
         yield modulestore
 
 
-class MixedModulestoreBuilder(object):
+class MixedModulestoreBuilder(StoreBuilderBase):
     """
     A builder class for a MixedModuleStore.
     """
@@ -210,7 +264,7 @@ class MixedModulestoreBuilder(object):
         self.mixed_modulestore = None
 
     @contextmanager
-    def build(self, contentstore):
+    def build_with_contentstore(self, contentstore):
         """
         A contextmanager that returns a mixed modulestore built on top of modulestores
         generated by other builder classes.
@@ -221,7 +275,7 @@ class MixedModulestoreBuilder(object):
         """
         names, generators = zip(*self.store_builders)
 
-        with nested(*(gen.build(contentstore) for gen in generators)) as modulestores:
+        with nested(*(gen.build_with_contentstore(contentstore) for gen in generators)) as modulestores:
             # Make the modulestore creation function just return the already-created modulestores
             store_iterator = iter(modulestores)
             create_modulestore_instance = lambda *args, **kwargs: store_iterator.next()
@@ -261,32 +315,6 @@ class MixedModulestoreBuilder(object):
             return store.db_connection.structures
 
 
-class MongoContentstoreBuilder(object):
-    """
-    A builder class for a MongoContentStore.
-    """
-    @contextmanager
-    def build(self):
-        """
-        A contextmanager that returns a MongoContentStore, and deletes its contents
-        when the context closes.
-        """
-        contentstore = MongoContentStore(
-            db='contentstore{}'.format(random.randint(0, 10000)),
-            collection='content',
-            **COMMON_DOCSTORE_CONFIG
-        )
-        contentstore.ensure_indexes()
-
-        try:
-            yield contentstore
-        finally:
-            # Delete the created database
-            contentstore._drop_database()
-
-    def __repr__(self):
-        return 'MongoContentstoreBuilder()'
-
 MIXED_MODULESTORE_BOTH_SETUP = MixedModulestoreBuilder([
     ('draft', MongoModulestoreBuilder()),
     ('split', VersioningModulestoreBuilder())
@@ -345,11 +373,11 @@ class CrossStoreXMLRoundtrip(CourseComparisonTest, PartitionTestCase):
         # Construct the contentstore for storing the first import
         with source_content_builder.build() as source_content:
             # Construct the modulestore for storing the first import (using the previously created contentstore)
-            with source_builder.build(source_content) as source_store:
+            with source_builder.build(contentstore=source_content) as source_store:
                 # Construct the contentstore for storing the second import
                 with dest_content_builder.build() as dest_content:
                     # Construct the modulestore for storing the second import (using the second contentstore)
-                    with dest_builder.build(dest_content) as dest_store:
+                    with dest_builder.build(contentstore=dest_content) as dest_store:
                         source_course_key = source_store.make_course_key('a', 'course', 'course')
                         dest_course_key = dest_store.make_course_key('a', 'course', 'course')
 

common/lib/xmodule/xmodule/modulestore/tests/test_mongo_call_count.py

+"""
+Tests to verify correct number of MongoDB calls during course import/export and traversal
+when using the Split modulestore.
+"""
+
+from tempfile import mkdtemp
+from shutil import rmtree
+from unittest import TestCase
+import ddt
+
+from xmodule.modulestore.xml_importer import import_from_xml
+from xmodule.modulestore.xml_exporter import export_to_xml
+from xmodule.modulestore.tests.factories import check_mongo_calls
+from xmodule.modulestore.tests.test_cross_modulestore_import_export import (
+    MixedModulestoreBuilder, VersioningModulestoreBuilder,
+    MongoModulestoreBuilder, TEST_DATA_DIR
+)
+
+MIXED_OLD_MONGO_MODULESTORE_BUILDER = MixedModulestoreBuilder([('draft', MongoModulestoreBuilder())])
+MIXED_SPLIT_MODULESTORE_BUILDER = MixedModulestoreBuilder([('split', VersioningModulestoreBuilder())])
+
+
+@ddt.ddt
+class CountMongoCallsXMLRoundtrip(TestCase):
+    """
+    This class exists to test XML import and export to/from Split.
+    """
+
+    def setUp(self):
+        super(CountMongoCallsXMLRoundtrip, self).setUp()
+        self.export_dir = mkdtemp()
+        self.addCleanup(rmtree, self.export_dir, ignore_errors=True)
+
+    @ddt.data(
+        (MIXED_OLD_MONGO_MODULESTORE_BUILDER, 287, 780, 702, 702),
+        (MIXED_SPLIT_MODULESTORE_BUILDER, 37, 16, 190, 189),
+    )
+    @ddt.unpack
+    def test_import_export(self, store_builder, export_reads, import_reads, first_import_writes, second_import_writes):
+        with store_builder.build() as (source_content, source_store):
+            with store_builder.build() as (dest_content, dest_store):
+                source_course_key = source_store.make_course_key('a', 'course', 'course')
+                dest_course_key = dest_store.make_course_key('a', 'course', 'course')
+
+                # An extra import write occurs in the first Split import due to the mismatch between
+                # the course id and the wiki_slug in the test XML course. The course must be updated
+                # with the correct wiki_slug during import.
+                with check_mongo_calls(import_reads, first_import_writes):
+                    import_from_xml(
+                        source_store,
+                        'test_user',
+                        TEST_DATA_DIR,
+                        course_dirs=['manual-testing-complete'],
+                        static_content_store=source_content,
+                        target_course_id=source_course_key,
+                        create_course_if_not_present=True,
+                        raise_on_failure=True,
+                    )
+
+                with check_mongo_calls(export_reads):
+                    export_to_xml(
+                        source_store,
+                        source_content,
+                        source_course_key,
+                        self.export_dir,
+                        'exported_source_course',
+                    )
+
+                with check_mongo_calls(import_reads, second_import_writes):
+                    import_from_xml(
+                        dest_store,
+                        'test_user',
+                        self.export_dir,
+                        course_dirs=['exported_source_course'],
+                        static_content_store=dest_content,
+                        target_course_id=dest_course_key,
+                        create_course_if_not_present=True,
+                        raise_on_failure=True,
+                    )
+
+
+@ddt.ddt
+class CountMongoCallsCourseTraversal(TestCase):
+    """
+    Tests the number of Mongo calls made when traversing a course tree from the top course root
+    to the leaf nodes.
+    """
+
+    @ddt.data(
+        (MIXED_OLD_MONGO_MODULESTORE_BUILDER, None, 189),  # The way this traversal *should* be done.
+        (MIXED_OLD_MONGO_MODULESTORE_BUILDER, 0, 387),     # The pathological case - do *not* query a course this way!
+        (MIXED_SPLIT_MODULESTORE_BUILDER, None, 7),  # The way this traversal *should* be done.
+        (MIXED_SPLIT_MODULESTORE_BUILDER, 0, 145)    # The pathological case - do *not* query a course this way!
+    )
+    @ddt.unpack
+    def test_number_mongo_calls(self, store, depth, num_mongo_calls):
+        with store.build() as (source_content, source_store):
+
+            source_course_key = source_store.make_course_key('a', 'course', 'course')
+
+            # First, import a course.
+            import_from_xml(
+                source_store,
+                'test_user',
+                TEST_DATA_DIR,
+                course_dirs=['manual-testing-complete'],
+                static_content_store=source_content,
+                target_course_id=source_course_key,
+                create_course_if_not_present=True,
+                raise_on_failure=True,
+            )
+
+            # Course traversal modeled after the traversal done here:
+            # lms/djangoapps/mobile_api/video_outlines/serializers.py:BlockOutline
+            # Starting at the root course block, do a breadth-first traversal using
+            # get_children() to retrieve each block's children.
+            with check_mongo_calls(num_mongo_calls):
+                start_block = source_store.get_course(source_course_key, depth=depth)
+                stack = [start_block]
+                while stack:
+                    curr_block = stack.pop()
+                    if curr_block.has_children:
+                        for block in reversed(curr_block.get_children()):
+                            stack.append(block)