PYTHON-3250 Speed up majority writes in test suite (#936)

Author: ShaneHarvey

test/__init__.py

@@ -1029,10 +1029,15 @@ def tearDown(self):
         super(MockClientTest, self).tearDown()
 
 
+# Global knobs to speed up the test suite.
+global_knobs = client_knobs(events_queue_frequency=0.05)
+
+
 def setup():
     client_context.init()
     warnings.resetwarnings()
     warnings.simplefilter("always")
+    global_knobs.enable()
 
 
 def _get_executors(topology):
@@ -1086,6 +1091,7 @@ def print_running_clients():
 
 
 def teardown():
+    global_knobs.disable()
     garbage = []
     for g in gc.garbage:
         garbage.append("GARBAGE: %r" % (g,))

test/test_encryption.py

@@ -24,7 +24,7 @@
 import textwrap
 import traceback
 import uuid
-from typing import Any
+from typing import Any, Dict
 
 from pymongo.collection import Collection
 
@@ -621,31 +621,30 @@ def maybe_skip_scenario(self, test):
     def setup_scenario(self, scenario_def):
         """Override a test's setup."""
         key_vault_data = scenario_def["key_vault_data"]
+        json_schema = scenario_def["json_schema"]
+        data = scenario_def["data"]
         if key_vault_data:
-            coll = client_context.client.get_database(
-                "keyvault", write_concern=WriteConcern(w="majority"), codec_options=OPTS
-            )["datakeys"]
-            coll.drop()
+            coll = client_context.client.get_database("keyvault", codec_options=OPTS)["datakeys"]
+            coll.delete_many({})
             coll.insert_many(key_vault_data)
 
         db_name = self.get_scenario_db_name(scenario_def)
         coll_name = self.get_scenario_coll_name(scenario_def)
-        db = client_context.client.get_database(
-            db_name, write_concern=WriteConcern(w="majority"), codec_options=OPTS
-        )
+        db = client_context.client.get_database(db_name, codec_options=OPTS)
         coll = db[coll_name]
         coll.drop()
-        json_schema = scenario_def["json_schema"]
+        wc = WriteConcern(w="majority")
+        kwargs: Dict[str, Any] = {}
         if json_schema:
-            db.create_collection(
-                coll_name, validator={"$jsonSchema": json_schema}, codec_options=OPTS
-            )
-        else:
-            db.create_collection(coll_name)
+            kwargs["validator"] = {"$jsonSchema": json_schema}
+            kwargs["codec_options"] = OPTS
+        if not data:
+            kwargs["write_concern"] = wc
+        db.create_collection(coll_name, **kwargs)
 
-        if scenario_def["data"]:
+        if data:
             # Load data.
-            coll.insert_many(scenario_def["data"])
+            coll.with_options(write_concern=wc).insert_many(scenario_def["data"])
 
     def allowable_errors(self, op):
         """Override expected error classes."""

test/test_retryable_reads.py

@@ -112,17 +112,20 @@ def get_scenario_coll_name(self, scenario_def):
     def setup_scenario(self, scenario_def):
         """Override a test's setup to support GridFS tests."""
         if "bucket_name" in scenario_def:
+            data = scenario_def["data"]
             db_name = self.get_scenario_db_name(scenario_def)
-            db = client_context.client.get_database(
-                db_name, write_concern=WriteConcern(w="majority")
-            )
-            # Create a bucket for the retryable reads GridFS tests.
-            client_context.client.drop_database(db_name)
-            if scenario_def["data"]:
-                data = scenario_def["data"]
-                # Load data.
+            db = client_context.client[db_name]
+            # Create a bucket for the retryable reads GridFS tests with as few
+            # majority writes as possible.
+            wc = WriteConcern(w="majority")
+            if data:
+                db["fs.chunks"].drop()
+                db["fs.files"].drop()
                 db["fs.chunks"].insert_many(data["fs.chunks"])
-                db["fs.files"].insert_many(data["fs.files"])
+                db.get_collection("fs.files", write_concern=wc).insert_many(data["fs.files"])
+            else:
+                db.get_collection("fs.chunks").drop()
+                db.get_collection("fs.files", write_concern=wc).drop()
         else:
             super(TestSpec, self).setup_scenario(scenario_def)
 

test/unified_format.py

@@ -721,22 +721,24 @@ def should_run_on(run_on_spec):
         return False
 
     def insert_initial_data(self, initial_data):
-        for collection_data in initial_data:
+        for i, collection_data in enumerate(initial_data):
             coll_name = collection_data["collectionName"]
             db_name = collection_data["databaseName"]
             documents = collection_data["documents"]
 
-            coll = self.client.get_database(db_name).get_collection(
-                coll_name, write_concern=WriteConcern(w="majority")
-            )
-            coll.drop()
-
-            if len(documents) > 0:
-                coll.insert_many(documents)
+            # Setup the collection with as few majority writes as possible.
+            db = self.client[db_name]
+            db.drop_collection(coll_name)
+            # Only use majority wc only on the final write.
+            if i == len(initial_data) - 1:
+                wc = WriteConcern(w="majority")
+            else:
+                wc = WriteConcern(w=1)
+            if documents:
+                db.get_collection(coll_name, write_concern=wc).insert_many(documents)
             else:
-                # ensure collection exists
-                result = coll.insert_one({})
-                coll.delete_one({"_id": result.inserted_id})
+                # Ensure collection exists
+                db.create_collection(coll_name, write_concern=wc)
 
     @classmethod
     def setUpClass(cls):

test/utils_spec_runner.py

@@ -453,13 +453,20 @@ def setup_scenario(self, scenario_def):
         """Allow specs to override a test's setup."""
         db_name = self.get_scenario_db_name(scenario_def)
         coll_name = self.get_scenario_coll_name(scenario_def)
-        db = client_context.client.get_database(db_name, write_concern=WriteConcern(w="majority"))
-        coll = db[coll_name]
-        coll.drop()
-        db.create_collection(coll_name)
-        if scenario_def["data"]:
-            # Load data.
-            coll.insert_many(scenario_def["data"])
+        documents = scenario_def["data"]
+
+        # Setup the collection with as few majority writes as possible.
+        db = client_context.client.get_database(db_name)
+        coll_exists = bool(db.list_collection_names(filter={"name": coll_name}))
+        if coll_exists:
+            db[coll_name].delete_many({})
+        # Only use majority wc only on the final write.
+        wc = WriteConcern(w="majority")
+        if documents:
+            db.get_collection(coll_name, write_concern=wc).insert_many(documents)
+        elif not coll_exists:
+            # Ensure collection exists.
+            db.create_collection(coll_name, write_concern=wc)
 
     def run_scenario(self, scenario_def, test):
         self.maybe_skip_scenario(test)