Merge pull request #14213 from edx/beryl/mongocrazy

Pre-load Definitions in Split Modulestore when accessing all blocks

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

@@ -432,9 +432,11 @@ class SplitBulkWriteMixin(BulkOperationsMixin):
 
         if len(ids):
             # Query the db for the definitions.
-            defs_from_db = self.db_connection.get_definitions(list(ids), course_key)
+            defs_from_db = list(self.db_connection.get_definitions(list(ids), course_key))
+            defs_dict = {d.get('_id'): d for d in defs_from_db}
             # Add the retrieved definitions to the cache.
-            bulk_write_record.definitions.update({d.get('_id'): d for d in defs_from_db})
+            bulk_write_record.definitions_in_db.update(defs_dict.iterkeys())
+            bulk_write_record.definitions.update(defs_dict)
             definitions.extend(defs_from_db)
         return definitions
 
@@ -772,11 +774,16 @@ class SplitMongoModuleStore(SplitBulkWriteMixin, ModuleStoreWriteBase):
         Load the definitions into each block if lazy is in kwargs and is False;
         otherwise, do not load the definitions - they'll be loaded later when needed.
         """
+        lazy = kwargs.pop('lazy', True)
+        should_cache_items = not lazy
+
         runtime = self._get_cache(course_entry.structure['_id'])
         if runtime is None:
-            lazy = kwargs.pop('lazy', True)
             runtime = self.create_runtime(course_entry, lazy)
             self._add_cache(course_entry.structure['_id'], runtime)
+            should_cache_items = True
+
+        if should_cache_items:
             self.cache_items(runtime, block_keys, course_entry.course_key, depth, lazy)
 
         return [runtime.load_item(block_key, course_entry, **kwargs) for block_key in block_keys]

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

@@ -13,7 +13,8 @@ from xmodule.modulestore.xml_exporter import export_course_to_xml
 from xmodule.modulestore.tests.factories import check_mongo_calls
 from xmodule.modulestore.tests.utils import (
     MixedModulestoreBuilder, VersioningModulestoreBuilder,
-    MongoModulestoreBuilder, TEST_DATA_DIR
+    MongoModulestoreBuilder, TEST_DATA_DIR,
+    MemoryCache,
 )
 
 MIXED_OLD_MONGO_MODULESTORE_BUILDER = MixedModulestoreBuilder([('draft', MongoModulestoreBuilder())])
@@ -87,6 +88,46 @@ class CountMongoCallsCourseTraversal(TestCase):
     to the leaf nodes.
     """
 
+    def _traverse_blocks_in_course(self, course, access_all_block_fields):
+        """
+        Traverses all the blocks in the given course.
+        If access_all_block_fields is True, also reads all the
+        xblock fields in each block in the course.
+        """
+        all_blocks = []
+        stack = [course]
+        while stack:
+            curr_block = stack.pop()
+            all_blocks.append(curr_block)
+            if curr_block.has_children:
+                for block in reversed(curr_block.get_children()):
+                    stack.append(block)
+
+        if access_all_block_fields:
+            # Read the fields on each block in order to ensure each block and its definition is loaded.
+            for xblock in all_blocks:
+                for __, field in xblock.fields.iteritems():
+                    if field.is_set_on(xblock):
+                        __ = field.read_from(xblock)
+
+    def _import_course(self, content_store, modulestore):
+        """
+        Imports a course for testing.
+        Returns the course key.
+        """
+        course_key = modulestore.make_course_key('a', 'course', 'course')
+        import_course_from_xml(
+            modulestore,
+            'test_user',
+            TEST_DATA_DIR,
+            source_dirs=['manual-testing-complete'],
+            static_content_store=content_store,
+            target_id=course_key,
+            create_if_not_present=True,
+            raise_on_failure=True,
+        )
+        return course_key
+
     # Suppose you want to traverse a course - maybe accessing the fields of each XBlock in the course,
     # maybe not. What parameters should one use for get_course() in order to minimize the number of
     # mongo calls? The tests below both ensure that code changes don't increase the number of mongo calls
@@ -109,52 +150,43 @@ class CountMongoCallsCourseTraversal(TestCase):
         # (if you'll eventually access all the fields and load all the definitions anyway).
         (MIXED_SPLIT_MODULESTORE_BUILDER, None, False, True, 4),
         (MIXED_SPLIT_MODULESTORE_BUILDER, None, True, True, 38),
-        (MIXED_SPLIT_MODULESTORE_BUILDER, 0, False, True, 131),
+        (MIXED_SPLIT_MODULESTORE_BUILDER, 0, False, True, 38),
         (MIXED_SPLIT_MODULESTORE_BUILDER, 0, True, True, 38),
         (MIXED_SPLIT_MODULESTORE_BUILDER, None, False, False, 4),
         (MIXED_SPLIT_MODULESTORE_BUILDER, None, True, False, 4),
-        # TODO: The call count below seems like a bug - should be 4?
-        # Seems to be related to using self.lazy in CachingDescriptorSystem.get_module_data().
-        (MIXED_SPLIT_MODULESTORE_BUILDER, 0, False, False, 131),
+        (MIXED_SPLIT_MODULESTORE_BUILDER, 0, False, False, 4),
         (MIXED_SPLIT_MODULESTORE_BUILDER, 0, True, False, 4)
     )
     @ddt.unpack
-    def test_number_mongo_calls(self, store, depth, lazy, access_all_block_fields, 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_course_from_xml(
-                source_store,
-                'test_user',
-                TEST_DATA_DIR,
-                source_dirs=['manual-testing-complete'],
-                static_content_store=source_content,
-                target_id=source_course_key,
-                create_if_not_present=True,
-                raise_on_failure=True,
-            )
+    def test_number_mongo_calls(self, store_builder, depth, lazy, access_all_block_fields, num_mongo_calls):
+        request_cache = MemoryCache()
+        with store_builder.build(request_cache=request_cache) as (content_store, modulestore):
+            course_key = self._import_course(content_store, modulestore)
 
             # 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):
-                with source_store.bulk_operations(source_course_key):
-                    start_block = source_store.get_course(source_course_key, depth=depth, lazy=lazy)
-                    all_blocks = []
-                    stack = [start_block]
-                    while stack:
-                        curr_block = stack.pop()
-                        all_blocks.append(curr_block)
-                        if curr_block.has_children:
-                            for block in reversed(curr_block.get_children()):
-                                stack.append(block)
-
-                    if access_all_block_fields:
-                        # Read the fields on each block in order to ensure each block and its definition is loaded.
-                        for xblock in all_blocks:
-                            for __, field in xblock.fields.iteritems():
-                                if field.is_set_on(xblock):
-                                    __ = field.read_from(xblock)
+                with modulestore.bulk_operations(course_key):
+                    start_block = modulestore.get_course(course_key, depth=depth, lazy=lazy)
+                    self._traverse_blocks_in_course(start_block, access_all_block_fields)
+
+    @ddt.data(
+        (MIXED_OLD_MONGO_MODULESTORE_BUILDER, 176),
+        (MIXED_SPLIT_MODULESTORE_BUILDER, 5),
+    )
+    @ddt.unpack
+    def test_lazy_when_course_previously_cached(self, store_builder, num_mongo_calls):
+        request_cache = MemoryCache()
+        with store_builder.build(request_cache=request_cache) as (content_store, modulestore):
+            course_key = self._import_course(content_store, modulestore)
+
+            with check_mongo_calls(num_mongo_calls):
+                with modulestore.bulk_operations(course_key):
+                    # assume the course was retrieved earlier
+                    course = modulestore.get_course(course_key, depth=0, lazy=True)
+
+                    # and then subsequently retrieved with the lazy and depth=None values
+                    course = modulestore.get_item(course.location, depth=None, lazy=False)
+                    self._traverse_blocks_in_course(course, access_all_block_fields=True)

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

+"""
+Tests for bulk operations in Split Modulestore.
+"""
+# pylint: disable=protected-access
 import copy
 import ddt
 import unittest
@@ -444,6 +448,17 @@ class TestBulkWriteMixinFindMethods(TestBulkWriteMixin):
             else:
                 self.assertNotIn(db_definition(_id), results)
 
+    def test_get_definitions_doesnt_update_db(self):
+        test_ids = [1, 2]
+        db_definition = lambda _id: {'db': 'definition', '_id': _id}
+
+        db_definitions = [db_definition(_id) for _id in test_ids]
+        self.conn.get_definitions.return_value = db_definitions
+        self.bulk._begin_bulk_operation(self.course_key)
+        self.bulk.get_definitions(self.course_key, test_ids)
+        self.bulk._end_bulk_operation(self.course_key)
+        self.assertFalse(self.conn.insert_definition.called)
+
     def test_no_bulk_find_structures_derived_from(self):
         ids = [Mock(name='id')]
         self.conn.find_structures_derived_from.return_value = [MagicMock(name='result')]

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

@@ -194,7 +194,7 @@ class MemoryCache(object):
     the modulestore, and stores the data in a dictionary in memory.
     """
     def __init__(self):
-        self._data = {}
+        self.data = {}
 
     def get(self, key, default=None):
         """
@@ -204,7 +204,7 @@ class MemoryCache(object):
             key: The key to update.
             default: The value to return if the key hasn't been set previously.
         """
-        return self._data.get(key, default)
+        return self.data.get(key, default)
 
     def set(self, key, value):
         """
@@ -214,7 +214,7 @@ class MemoryCache(object):
             key: The key to update.
             value: The value change the key to.
         """
-        self._data[key] = value
+        self.data[key] = value
 
 
 class MongoContentstoreBuilder(object):
@@ -255,19 +255,19 @@ class StoreBuilderBase(object):
         """
         contentstore = kwargs.pop('contentstore', None)
         if not contentstore:
-            with self.build_without_contentstore() as (contentstore, modulestore):
+            with self.build_without_contentstore(**kwargs) as (contentstore, modulestore):
                 yield contentstore, modulestore
         else:
-            with self.build_with_contentstore(contentstore) as modulestore:
+            with self.build_with_contentstore(contentstore, **kwargs) as modulestore:
                 yield modulestore
 
     @contextmanager
-    def build_without_contentstore(self):
+    def build_without_contentstore(self, **kwargs):
         """
         Build both the contentstore and the modulestore.
         """
         with MongoContentstoreBuilder().build() as contentstore:
-            with self.build_with_contentstore(contentstore) as modulestore:
+            with self.build_with_contentstore(contentstore, **kwargs) as modulestore:
                 yield contentstore, modulestore
 
 
@@ -276,7 +276,7 @@ class MongoModulestoreBuilder(StoreBuilderBase):
     A builder class for a DraftModuleStore.
     """
     @contextmanager
-    def build_with_contentstore(self, contentstore):
+    def build_with_contentstore(self, contentstore, **kwargs):
         """
         A contextmanager that returns an isolated mongo modulestore, and then deletes
         all of its data at the end of the context.
@@ -324,7 +324,7 @@ class VersioningModulestoreBuilder(StoreBuilderBase):
     A builder class for a VersioningModuleStore.
     """
     @contextmanager
-    def build_with_contentstore(self, contentstore):
+    def build_with_contentstore(self, contentstore, **kwargs):
         """
         A contextmanager that returns an isolated versioning modulestore, and then deletes
         all of its data at the end of the context.
@@ -347,6 +347,7 @@ class VersioningModulestoreBuilder(StoreBuilderBase):
             fs_root,
             render_template=repr,
             xblock_mixins=XBLOCK_MIXINS,
+            **kwargs
         )
         modulestore.ensure_indexes()
 
@@ -369,7 +370,7 @@ class XmlModulestoreBuilder(StoreBuilderBase):
     """
     # pylint: disable=unused-argument
     @contextmanager
-    def build_with_contentstore(self, contentstore=None, course_ids=None):
+    def build_with_contentstore(self, contentstore=None, course_ids=None, **kwargs):
         """
         A contextmanager that returns an isolated xml modulestore
 
@@ -403,7 +404,7 @@ class MixedModulestoreBuilder(StoreBuilderBase):
         self.mixed_modulestore = None
 
     @contextmanager
-    def build_with_contentstore(self, contentstore):
+    def build_with_contentstore(self, contentstore, **kwargs):
         """
         A contextmanager that returns a mixed modulestore built on top of modulestores
         generated by other builder classes.
@@ -414,7 +415,7 @@ class MixedModulestoreBuilder(StoreBuilderBase):
         """
         names, generators = zip(*self.store_builders)
 
-        with nested(*(gen.build_with_contentstore(contentstore) for gen in generators)) as modulestores:
+        with nested(*(gen.build_with_contentstore(contentstore, **kwargs) for gen in generators)) as modulestores:
             # Make the modulestore creation function just return the already-created modulestores
             store_iterator = iter(modulestores)
             next_modulestore = lambda *args, **kwargs: store_iterator.next()

lms/djangoapps/course_api/blocks/tests/test_api.py

@@ -2,12 +2,14 @@
 Tests for Blocks api.py
 """
 
+import ddt
 from django.test.client import RequestFactory
 
+from openedx.core.djangoapps.content.block_structure.api import clear_course_from_cache
 from student.tests.factories import UserFactory
 from xmodule.modulestore import ModuleStoreEnum
 from xmodule.modulestore.tests.django_utils import SharedModuleStoreTestCase
-from xmodule.modulestore.tests.factories import SampleCourseFactory
+from xmodule.modulestore.tests.factories import SampleCourseFactory, check_mongo_calls
 
 from ..api import get_blocks
 
@@ -19,7 +21,8 @@ class TestGetBlocks(SharedModuleStoreTestCase):
     @classmethod
     def setUpClass(cls):
         super(TestGetBlocks, cls).setUpClass()
-        cls.course = SampleCourseFactory.create()
+        with cls.store.default_store(ModuleStoreEnum.Type.split):
+            cls.course = SampleCourseFactory.create()
 
         # hide the html block
         cls.html_block = cls.store.get_item(cls.course.id.make_usage_key('html', 'html_x1a_1'))
@@ -95,3 +98,47 @@ class TestGetBlocks(SharedModuleStoreTestCase):
         self.assertEquals(len(blocks['blocks']), 3)
         for block in blocks['blocks'].itervalues():
             self.assertEqual(block['type'], 'problem')
+
+
+@ddt.ddt
+class TestGetBlocksQueryCounts(SharedModuleStoreTestCase):
+    """
+    Tests query counts for the get_blocks function.
+    """
+    def setUp(self):
+        super(TestGetBlocksQueryCounts, self).setUp()
+
+        self.user = UserFactory.create()
+        self.request = RequestFactory().get("/dummy")
+        self.request.user = self.user
+
+    def _create_course(self, store_type):
+        """
+        Creates the sample course in the given store type.
+        """
+        with self.store.default_store(store_type):
+            return SampleCourseFactory.create()
+
+    def _get_blocks(self, course, expected_mongo_queries):
+        """
+        Verifies the number of expected queries when calling
+        get_blocks on the given course.
+        """
+        with check_mongo_calls(expected_mongo_queries):
+            with self.assertNumQueries(2):
+                get_blocks(self.request, course.location, self.user)
+
+    @ddt.data(ModuleStoreEnum.Type.mongo, ModuleStoreEnum.Type.split)
+    def test_query_counts_cached(self, store_type):
+        course = self._create_course(store_type)
+        self._get_blocks(course, expected_mongo_queries=0)
+
+    @ddt.data(
+        (ModuleStoreEnum.Type.mongo, 5),
+        (ModuleStoreEnum.Type.split, 3),
+    )
+    @ddt.unpack
+    def test_query_counts_uncached(self, store_type, expected_mongo_queries):
+        course = self._create_course(store_type)
+        clear_course_from_cache(course.id)
+        self._get_blocks(course, expected_mongo_queries)

lms/djangoapps/courseware/tests/test_views.py

@@ -191,6 +191,48 @@ class TestJumpTo(ModuleStoreTestCase):
 
 @attr(shard=2)
 @ddt.ddt
+class IndexQueryTestCase(ModuleStoreTestCase):
+    """
+    Tests for query count.
+    """
+    CREATE_USER = False
+    NUM_PROBLEMS = 20
+
+    @ddt.data(
+        (ModuleStoreEnum.Type.mongo, 8),
+        (ModuleStoreEnum.Type.split, 4),
+    )
+    @ddt.unpack
+    def test_index_query_counts(self, store_type, expected_query_count):
+        with self.store.default_store(store_type):
+            course = CourseFactory.create()
+            with self.store.bulk_operations(course.id):
+                chapter = ItemFactory.create(category='chapter', parent_location=course.location)
+                section = ItemFactory.create(category='sequential', parent_location=chapter.location)
+                vertical = ItemFactory.create(category='vertical', parent_location=section.location)
+                for _ in range(self.NUM_PROBLEMS):
+                    ItemFactory.create(category='problem', parent_location=vertical.location)
+
+        password = 'test'
+        self.user = UserFactory(password=password)
+        self.client.login(username=self.user.username, password=password)
+        CourseEnrollment.enroll(self.user, course.id)
+
+        with check_mongo_calls(expected_query_count):
+            url = reverse(
+                'courseware_section',
+                kwargs={
+                    'course_id': unicode(course.id),
+                    'chapter': unicode(chapter.location.name),
+                    'section': unicode(section.location.name),
+                }
+            )
+            response = self.client.get(url)
+            self.assertEqual(response.status_code, 200)
+
+
+@attr(shard=2)
+@ddt.ddt
 class ViewsTestCase(ModuleStoreTestCase):
     """
     Tests for views.py methods.

lms/djangoapps/courseware/views/index.py

@@ -348,7 +348,7 @@ class CoursewareIndex(View):
         sets up the runtime, which binds the request user to the section.
         """
         # Pre-fetch all descendant data
-        self.section = modulestore().get_item(self.section.location, depth=None)
+        self.section = modulestore().get_item(self.section.location, depth=None, lazy=False)
         self.field_data_cache.add_descriptor_descendents(self.section, depth=None)
 
         # Bind section to user

lms/djangoapps/grades/tasks.py

@@ -158,43 +158,45 @@ def _update_subsection_grades(
     that those subsection grades were updated.
     """
     student = User.objects.get(id=user_id)
-    course_structure = get_course_blocks(student, modulestore().make_course_usage_key(course_key))
-    subsections_to_update = course_structure.get_transformer_block_field(
-        scored_block_usage_key,
-        GradesTransformer,
-        'subsections',
-        set(),
-    )
-
-    course = modulestore().get_course(course_key, depth=0)
-    subsection_grade_factory = SubsectionGradeFactory(student, course, course_structure)
-
-    try:
-        for subsection_usage_key in subsections_to_update:
-            if subsection_usage_key in course_structure:
-                subsection_grade = subsection_grade_factory.update(
-                    course_structure[subsection_usage_key],
-                    only_if_higher,
-                )
-                SUBSECTION_SCORE_CHANGED.send(
-                    sender=recalculate_subsection_grade,
-                    course=course,
-                    course_structure=course_structure,
-                    user=student,
-                    subsection_grade=subsection_grade,
-                )
-
-    except DatabaseError as exc:
-        raise _retry_recalculate_subsection_grade(
-            user_id,
-            course_id,
-            usage_id,
-            only_if_higher,
-            expected_modified_time,
-            score_deleted,
-            exc,
+    store = modulestore()
+    with store.bulk_operations(course_key):
+        course_structure = get_course_blocks(student, store.make_course_usage_key(course_key))
+        subsections_to_update = course_structure.get_transformer_block_field(
+            scored_block_usage_key,
+            GradesTransformer,
+            'subsections',
+            set(),
         )
 
+        course = store.get_course(course_key, depth=0)
+        subsection_grade_factory = SubsectionGradeFactory(student, course, course_structure)
+
+        try:
+            for subsection_usage_key in subsections_to_update:
+                if subsection_usage_key in course_structure:
+                    subsection_grade = subsection_grade_factory.update(
+                        course_structure[subsection_usage_key],
+                        only_if_higher,
+                    )
+                    SUBSECTION_SCORE_CHANGED.send(
+                        sender=recalculate_subsection_grade,
+                        course=course,
+                        course_structure=course_structure,
+                        user=student,
+                        subsection_grade=subsection_grade,
+                    )
+
+        except DatabaseError as exc:
+            raise _retry_recalculate_subsection_grade(
+                user_id,
+                course_id,
+                usage_id,
+                only_if_higher,
+                expected_modified_time,
+                score_deleted,
+                exc,
+            )
+
 
 def _retry_recalculate_subsection_grade(
         user_id,

lms/djangoapps/grades/tests/test_tasks.py

@@ -121,16 +121,17 @@ class RecalculateSubsectionGradeTest(ModuleStoreTestCase):
         self.assertTrue(mock_subsection_signal.called)
 
     @ddt.data(
-        (ModuleStoreEnum.Type.mongo, 1),
-        (ModuleStoreEnum.Type.split, 0),
+        (ModuleStoreEnum.Type.mongo, 1, 23),
+        (ModuleStoreEnum.Type.split, 3, 22),
     )
     @ddt.unpack
-    def test_subsection_grade_updated(self, default_store, added_queries):
+    def test_subsection_grade_updated(self, default_store, num_mongo_calls, num_sql_calls):
         with self.store.default_store(default_store):
             self.set_up_course()
             self.assertTrue(PersistentGradesEnabledFlag.feature_enabled(self.course.id))
-            with check_mongo_calls(2) and self.assertNumQueries(22 + added_queries):
-                self._apply_recalculate_subsection_grade()
+            with check_mongo_calls(num_mongo_calls):
+                with self.assertNumQueries(num_sql_calls):
+                    self._apply_recalculate_subsection_grade()
 
     @patch('lms.djangoapps.grades.signals.signals.SUBSECTION_SCORE_CHANGED.send')
     def test_other_inaccessible_subsection(self, mock_subsection_signal):
@@ -166,27 +167,38 @@ class RecalculateSubsectionGradeTest(ModuleStoreTestCase):
             self._apply_recalculate_subsection_grade()
             self.assertEquals(mock_block_structure_create.call_count, 1)
 
+    # TODO (TNL-6225) Fix the number of SQL queries so they
+    # don't grow linearly with the number of sequentials.
     @ddt.data(
-        (ModuleStoreEnum.Type.mongo, 1),
-        (ModuleStoreEnum.Type.split, 0),
+        (ModuleStoreEnum.Type.mongo, 1, 46),
+        (ModuleStoreEnum.Type.split, 3, 45),
     )
     @ddt.unpack
-    def test_query_count_does_not_change_with_more_problems(self, default_store, added_queries):
+    def test_query_count_does_not_change_with_more_content(self, default_store, num_mongo_calls, num_sql_calls):
         with self.store.default_store(default_store):
             self.set_up_course()
             self.assertTrue(PersistentGradesEnabledFlag.feature_enabled(self.course.id))
-            ItemFactory.create(parent=self.sequential, category='problem', display_name='problem2')
-            ItemFactory.create(parent=self.sequential, category='problem', display_name='problem3')
-            with check_mongo_calls(2) and self.assertNumQueries(22 + added_queries):
-                self._apply_recalculate_subsection_grade()
+
+            num_problems = 10
+            for _ in range(num_problems):
+                ItemFactory.create(parent=self.sequential, category='problem')
+
+            num_sequentials = 10
+            for _ in range(num_sequentials):
+                ItemFactory.create(parent=self.chapter, category='sequential')
+
+            with check_mongo_calls(num_mongo_calls):
+                with self.assertNumQueries(num_sql_calls):
+                    self._apply_recalculate_subsection_grade()
 
     @ddt.data(ModuleStoreEnum.Type.mongo, ModuleStoreEnum.Type.split)
     def test_subsection_grades_not_enabled_on_course(self, default_store):
         with self.store.default_store(default_store):
             self.set_up_course(enable_subsection_grades=False)
             self.assertFalse(PersistentGradesEnabledFlag.feature_enabled(self.course.id))
-            with check_mongo_calls(2) and self.assertNumQueries(0):
-                self._apply_recalculate_subsection_grade()
+            with check_mongo_calls(0):
+                with self.assertNumQueries(0):
+                    self._apply_recalculate_subsection_grade()
 
     @skip("Pending completion of TNL-5089")
     @ddt.data(
@@ -200,8 +212,9 @@ class RecalculateSubsectionGradeTest(ModuleStoreTestCase):
         PersistentGradesEnabledFlag.objects.create(enabled=feature_flag)
         with self.store.default_store(default_store):
             self.set_up_course()
-            with check_mongo_calls(0) and self.assertNumQueries(3 if feature_flag else 2):
-                self._apply_recalculate_subsection_grade()
+            with check_mongo_calls(0):
+                with self.assertNumQueries(3 if feature_flag else 2):
+                    self._apply_recalculate_subsection_grade()
 
     @patch('lms.djangoapps.grades.tasks.recalculate_subsection_grade_v2.retry')
     @patch('lms.djangoapps.grades.new.subsection_grade.SubsectionGradeFactory.update')

lms/djangoapps/grades/tests/test_transformer.py

@@ -17,7 +17,7 @@ from xmodule.modulestore.tests.factories import check_mongo_calls
 
 from lms.djangoapps.course_blocks.api import get_course_blocks
 from lms.djangoapps.course_blocks.transformers.tests.helpers import CourseStructureTestCase
-from openedx.core.djangoapps.content.block_structure.api import get_cache
+from openedx.core.djangoapps.content.block_structure.api import clear_course_from_cache
 from ..transformer import GradesTransformer
 
 
@@ -390,6 +390,7 @@ class GradesTransformerTestCase(CourseStructureTestCase):
         )
 
 
+@ddt.ddt
 class MultiProblemModulestoreAccessTestCase(CourseStructureTestCase, SharedModuleStoreTestCase):
     """
     Test mongo usage in GradesTransformer.
@@ -403,7 +404,12 @@ class MultiProblemModulestoreAccessTestCase(CourseStructureTestCase, SharedModul
         self.student = UserFactory.create(is_staff=False, username=u'test_student', password=password)
         self.client.login(username=self.student.username, password=password)
 
-    def test_modulestore_performance(self):
+    @ddt.data(
+        (ModuleStoreEnum.Type.split, 3),
+        (ModuleStoreEnum.Type.mongo, 2),
+    )
+    @ddt.unpack
+    def test_modulestore_performance(self, store_type, expected_mongo_queries):
         """
         Test that a constant number of mongo calls are made regardless of how
         many grade-related blocks are in the course.
@@ -437,7 +443,8 @@ class MultiProblemModulestoreAccessTestCase(CourseStructureTestCase, SharedModul
                         </problem>'''.format(number=problem_number),
                 }
             )
-        blocks = self.build_course(course)
-        get_cache().clear()
-        with check_mongo_calls(2):
+        with self.store.default_store(store_type):
+            blocks = self.build_course(course)
+        clear_course_from_cache(blocks[u'course'].id)
+        with check_mongo_calls(expected_mongo_queries):
             get_course_blocks(self.student, blocks[u'course'].location, self.transformers)

openedx/core/djangoapps/bookmarks/tests/test_tasks.py

@@ -103,11 +103,11 @@ class XBlockCacheTaskTests(BookmarksTestsBase):
         }
 
     @ddt.data(
-        (ModuleStoreEnum.Type.mongo, 2, 2, 3),
-        (ModuleStoreEnum.Type.mongo, 4, 2, 3),
-        (ModuleStoreEnum.Type.mongo, 2, 3, 4),
-        (ModuleStoreEnum.Type.mongo, 4, 3, 4),
-        (ModuleStoreEnum.Type.mongo, 2, 4, 5),
+        (ModuleStoreEnum.Type.mongo, 2, 2, 4),
+        (ModuleStoreEnum.Type.mongo, 4, 2, 4),
+        (ModuleStoreEnum.Type.mongo, 2, 3, 5),
+        (ModuleStoreEnum.Type.mongo, 4, 3, 5),
+        (ModuleStoreEnum.Type.mongo, 2, 4, 6),
         # (ModuleStoreEnum.Type.mongo, 4, 4, 6), Too slow.
         (ModuleStoreEnum.Type.split, 2, 2, 3),
         (ModuleStoreEnum.Type.split, 4, 2, 3),
@@ -119,6 +119,9 @@ class XBlockCacheTaskTests(BookmarksTestsBase):
 
         course = self.create_course_with_blocks(children_per_block, depth, store_type)
 
+        # clear cache to get consistent query counts
+        self.clear_caches()
+
         with check_mongo_calls(expected_mongo_calls):
             blocks_data = _calculate_course_xblocks_data(course.id)
             self.assertGreater(len(blocks_data), children_per_block ** depth)

openedx/core/lib/block_structure/factory.py

@@ -53,7 +53,7 @@ class BlockStructureFactory(object):
                 block_structure._add_relation(xblock.location, child.location)  # pylint: disable=protected-access
                 build_block_structure(child)
 
-        root_xblock = modulestore.get_item(root_block_usage_key, depth=None)
+        root_xblock = modulestore.get_item(root_block_usage_key, depth=None, lazy=False)
         build_block_structure(root_xblock)
         return block_structure
 

openedx/core/lib/block_structure/tests/helpers.py

@@ -55,7 +55,7 @@ class MockModulestore(object):
         """
         self.blocks = blocks
 
-    def get_item(self, block_key, depth=None):  # pylint: disable=unused-argument
+    def get_item(self, block_key, depth=None, lazy=False):  # pylint: disable=unused-argument
         """
         Returns the mock XBlock (MockXBlock) associated with the
         given block_key.