Count & compare db hits between old and split mongo

LMS-11090

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

@@ -9,7 +9,7 @@ from xblock.core import XBlock
 from xmodule.tabs import StaticTab
 from decorator import contextmanager
 from mock import Mock, patch
-from nose.tools import assert_less_equal, assert_greater_equal
+from nose.tools import assert_less_equal, assert_greater_equal, assert_equal
 
 
 class Dummy(object):
@@ -259,18 +259,49 @@ def check_mongo_calls(mongo_store, num_finds=0, num_sends=None):
     of calls to send_message which is for insert, update, and remove (if you provide num_sends). At the
     end of the with statement, it compares the counts to the num_finds and num_sends.
 
-    :param mongo_store: the MongoModulestore or subclass to watch
+    :param mongo_store: the MongoModulestore or subclass to watch or a SplitMongoModuleStore
     :param num_finds: the exact number of find calls expected
     :param num_sends: If none, don't instrument the send calls. If non-none, count and compare to
         the given int value.
     """
-    with check_exact_number_of_calls(mongo_store.collection, mongo_store.collection.find, num_finds):
-        if num_sends:
-            with check_exact_number_of_calls(
-                mongo_store.database.connection,
-                mongo_store.database.connection._send_message,  # pylint: disable=protected-access
-                num_sends,
-            ):
+    if mongo_store.get_modulestore_type() == ModuleStoreEnum.Type.mongo:
+        with check_exact_number_of_calls(mongo_store.collection, mongo_store.collection.find, num_finds):
+            if num_sends:
+                with check_exact_number_of_calls(
+                    mongo_store.database.connection,
+                    mongo_store.database.connection._send_message,  # pylint: disable=protected-access
+                    num_sends,
+                ):
+                    yield
+            else:
                 yield
-        else:
-            yield
+    elif mongo_store.get_modulestore_type() == ModuleStoreEnum.Type.split:
+        collections = [
+            mongo_store.db_connection.course_index,
+            mongo_store.db_connection.structures,
+            mongo_store.db_connection.definitions,
+        ]
+        # could add else clause which raises exception or just rely on the below to suss that out
+        try:
+            find_wraps = []
+            wrap_patches = []
+            for collection in collections:
+                find_wrap = Mock(wraps=collection.find)
+                find_wraps.append(find_wrap)
+                wrap_patch = patch.object(collection, 'find', find_wrap)
+                wrap_patches.append(wrap_patch)
+                wrap_patch.start()
+            if num_sends:
+                connection = mongo_store.db_connection.database.connection
+                with check_exact_number_of_calls(
+                    connection,
+                    connection._send_message,  # pylint: disable=protected-access
+                    num_sends,
+                ):
+                    yield
+            else:
+                yield
+        finally:
+            map(lambda wrap_patch: wrap_patch.stop(), wrap_patches)
+            call_count = sum([find_wrap.call_count for find_wrap in find_wraps])
+            assert_equal(call_count, num_finds)