PYTHON-4650 Fix MongoClient.bulk_write test failure when compression is enabled (mongodb#1786)

Author: shruti-sridhar

test/__init__.py

@@ -852,6 +852,10 @@ def max_bson_size(self):
     def max_write_batch_size(self):
         return (self.hello)["maxWriteBatchSize"]
 
+    @property
+    def max_message_size_bytes(self):
+        return (self.hello)["maxMessageSizeBytes"]
+
 
 # Reusable client context
 client_context = ClientContext()

test/asynchronous/__init__.py

@@ -854,6 +854,10 @@ async def max_bson_size(self):
     async def max_write_batch_size(self):
         return (await self.hello)["maxWriteBatchSize"]
 
+    @property
+    async def max_message_size_bytes(self):
+        return (await self.hello)["maxMessageSizeBytes"]
+
 
 # Reusable client context
 async_client_context = AsyncClientContext()

test/asynchronous/test_client_bulk_write.py

@@ -56,20 +56,24 @@ async def test_returns_error_if_no_namespace_provided(self):
 
 # https://github.com/mongodb/specifications/tree/master/source/crud/tests
 class TestClientBulkWriteCRUD(AsyncIntegrationTest):
+    async def asyncSetUp(self):
+        self.max_write_batch_size = await async_client_context.max_write_batch_size
+        self.max_bson_object_size = await async_client_context.max_bson_size
+        self.max_message_size_bytes = await async_client_context.max_message_size_bytes
+
     @async_client_context.require_version_min(8, 0, 0, -24)
     async def test_batch_splits_if_num_operations_too_large(self):
         listener = OvertCommandListener()
         client = await async_rs_or_single_client(event_listeners=[listener])
         self.addAsyncCleanup(client.aclose)
 
-        max_write_batch_size = (await async_client_context.hello)["maxWriteBatchSize"]
         models = []
-        for _ in range(max_write_batch_size + 1):
+        for _ in range(self.max_write_batch_size + 1):
             models.append(InsertOne(namespace="db.coll", document={"a": "b"}))
         self.addAsyncCleanup(client.db["coll"].drop)
 
         result = await client.bulk_write(models=models)
-        self.assertEqual(result.inserted_count, max_write_batch_size + 1)
+        self.assertEqual(result.inserted_count, self.max_write_batch_size + 1)
 
         bulk_write_events = []
         for event in listener.started_events:
@@ -78,7 +82,7 @@ async def test_batch_splits_if_num_operations_too_large(self):
         self.assertEqual(len(bulk_write_events), 2)
 
         first_event, second_event = bulk_write_events
-        self.assertEqual(len(first_event.command["ops"]), max_write_batch_size)
+        self.assertEqual(len(first_event.command["ops"]), self.max_write_batch_size)
         self.assertEqual(len(second_event.command["ops"]), 1)
         self.assertEqual(first_event.operation_id, second_event.operation_id)
 
@@ -88,12 +92,9 @@ async def test_batch_splits_if_ops_payload_too_large(self):
         client = await async_rs_or_single_client(event_listeners=[listener])
         self.addAsyncCleanup(client.aclose)
 
-        max_message_size_bytes = (await async_client_context.hello)["maxMessageSizeBytes"]
-        max_bson_object_size = (await async_client_context.hello)["maxBsonObjectSize"]
-
         models = []
-        num_models = int(max_message_size_bytes / max_bson_object_size + 1)
-        b_repeated = "b" * (max_bson_object_size - 500)
+        num_models = int(self.max_message_size_bytes / self.max_bson_object_size + 1)
+        b_repeated = "b" * (self.max_bson_object_size - 500)
         for _ in range(num_models):
             models.append(
                 InsertOne(
@@ -126,7 +127,6 @@ async def test_collects_write_concern_errors_across_batches(self):
             retryWrites=False,
         )
         self.addAsyncCleanup(client.aclose)
-        max_write_batch_size = (await async_client_context.hello)["maxWriteBatchSize"]
 
         fail_command = {
             "configureFailPoint": "failCommand",
@@ -138,7 +138,7 @@ async def test_collects_write_concern_errors_across_batches(self):
         }
         async with self.fail_point(fail_command):
             models = []
-            for _ in range(max_write_batch_size + 1):
+            for _ in range(self.max_write_batch_size + 1):
                 models.append(
                     InsertOne(
                         namespace="db.coll",
@@ -152,7 +152,7 @@ async def test_collects_write_concern_errors_across_batches(self):
             self.assertEqual(len(context.exception.write_concern_errors), 2)  # type: ignore[arg-type]
             self.assertIsNotNone(context.exception.partial_result)
             self.assertEqual(
-                context.exception.partial_result.inserted_count, max_write_batch_size + 1
+                context.exception.partial_result.inserted_count, self.max_write_batch_size + 1
             )
 
         bulk_write_events = []
@@ -172,9 +172,8 @@ async def test_collects_write_errors_across_batches_unordered(self):
         await collection.drop()
         await collection.insert_one(document={"_id": 1})
 
-        max_write_batch_size = (await async_client_context.hello)["maxWriteBatchSize"]
         models = []
-        for _ in range(max_write_batch_size + 1):
+        for _ in range(self.max_write_batch_size + 1):
             models.append(
                 InsertOne(
                     namespace="db.coll",
@@ -184,7 +183,7 @@ async def test_collects_write_errors_across_batches_unordered(self):
 
         with self.assertRaises(ClientBulkWriteException) as context:
             await client.bulk_write(models=models, ordered=False)
-        self.assertEqual(len(context.exception.write_errors), max_write_batch_size + 1)  # type: ignore[arg-type]
+        self.assertEqual(len(context.exception.write_errors), self.max_write_batch_size + 1)  # type: ignore[arg-type]
 
         bulk_write_events = []
         for event in listener.started_events:
@@ -203,9 +202,8 @@ async def test_collects_write_errors_across_batches_ordered(self):
         await collection.drop()
         await collection.insert_one(document={"_id": 1})
 
-        max_write_batch_size = (await async_client_context.hello)["maxWriteBatchSize"]
         models = []
-        for _ in range(max_write_batch_size + 1):
+        for _ in range(self.max_write_batch_size + 1):
             models.append(
                 InsertOne(
                     namespace="db.coll",
@@ -233,10 +231,9 @@ async def test_handles_cursor_requiring_getMore(self):
         self.addAsyncCleanup(collection.drop)
         await collection.drop()
 
-        max_bson_object_size = (await async_client_context.hello)["maxBsonObjectSize"]
         models = []
-        a_repeated = "a" * (max_bson_object_size // 2)
-        b_repeated = "b" * (max_bson_object_size // 2)
+        a_repeated = "a" * (self.max_bson_object_size // 2)
+        b_repeated = "b" * (self.max_bson_object_size // 2)
         models.append(
             UpdateOne(
                 namespace="db.coll",
@@ -275,12 +272,11 @@ async def test_handles_cursor_requiring_getMore_within_transaction(self):
         self.addAsyncCleanup(collection.drop)
         await collection.drop()
 
-        max_bson_object_size = (await async_client_context.hello)["maxBsonObjectSize"]
         async with client.start_session() as session:
             await session.start_transaction()
             models = []
-            a_repeated = "a" * (max_bson_object_size // 2)
-            b_repeated = "b" * (max_bson_object_size // 2)
+            a_repeated = "a" * (self.max_bson_object_size // 2)
+            b_repeated = "b" * (self.max_bson_object_size // 2)
             models.append(
                 UpdateOne(
                     namespace="db.coll",
@@ -319,16 +315,15 @@ async def test_handles_getMore_error(self):
         self.addAsyncCleanup(collection.drop)
         await collection.drop()
 
-        max_bson_object_size = (await async_client_context.hello)["maxBsonObjectSize"]
         fail_command = {
             "configureFailPoint": "failCommand",
             "mode": {"times": 1},
             "data": {"failCommands": ["getMore"], "errorCode": 8},
         }
         async with self.fail_point(fail_command):
             models = []
-            a_repeated = "a" * (max_bson_object_size // 2)
-            b_repeated = "b" * (max_bson_object_size // 2)
+            a_repeated = "a" * (self.max_bson_object_size // 2)
+            b_repeated = "b" * (self.max_bson_object_size // 2)
             models.append(
                 UpdateOne(
                     namespace="db.coll",
@@ -370,8 +365,7 @@ async def test_returns_error_if_unacknowledged_too_large_insert(self):
         client = await async_rs_or_single_client(event_listeners=[listener])
         self.addAsyncCleanup(client.aclose)
 
-        max_bson_object_size = (await async_client_context.hello)["maxBsonObjectSize"]
-        b_repeated = "b" * max_bson_object_size
+        b_repeated = "b" * self.max_bson_object_size
 
         # Insert document.
         models_insert = [InsertOne(namespace="db.coll", document={"a": b_repeated})]
@@ -384,25 +378,35 @@ async def test_returns_error_if_unacknowledged_too_large_insert(self):
             await client.bulk_write(models=models_replace, write_concern=WriteConcern(w=0))
 
     async def _setup_namespace_test_models(self):
-        max_message_size_bytes = (await async_client_context.hello)["maxMessageSizeBytes"]
-        max_bson_object_size = (await async_client_context.hello)["maxBsonObjectSize"]
-
-        ops_bytes = max_message_size_bytes - 1122
-        num_models = ops_bytes // max_bson_object_size
-        remainder_bytes = ops_bytes % max_bson_object_size
+        # See prose test specification below for details on these calculations.
+        # https://github.com/mongodb/specifications/tree/master/source/crud/tests#details-on-size-calculations
+        _EXISTING_BULK_WRITE_BYTES = 1122
+        _OPERATION_DOC_BYTES = 57
+        _NAMESPACE_DOC_BYTES = 217
+
+        # When compression is enabled, max_message_size is
+        # smaller to account for compression message header.
+        if async_client_context.client_options.get("compressors"):
+            max_message_size_bytes = self.max_message_size_bytes - 16
+        else:
+            max_message_size_bytes = self.max_message_size_bytes
+
+        ops_bytes = max_message_size_bytes - _EXISTING_BULK_WRITE_BYTES
+        num_models = ops_bytes // self.max_bson_object_size
+        remainder_bytes = ops_bytes % self.max_bson_object_size
 
         models = []
-        b_repeated = "b" * (max_bson_object_size - 57)
+        b_repeated = "b" * (self.max_bson_object_size - _OPERATION_DOC_BYTES)
         for _ in range(num_models):
             models.append(
                 InsertOne(
                     namespace="db.coll",
                     document={"a": b_repeated},
                 )
             )
-        if remainder_bytes >= 217:
+        if remainder_bytes >= _NAMESPACE_DOC_BYTES:
             num_models += 1
-            b_repeated = "b" * (remainder_bytes - 57)
+            b_repeated = "b" * (remainder_bytes - _OPERATION_DOC_BYTES)
             models.append(
                 InsertOne(
                     namespace="db.coll",
@@ -485,17 +489,15 @@ async def test_returns_error_if_no_writes_can_be_added_to_ops(self):
         client = await async_rs_or_single_client()
         self.addAsyncCleanup(client.aclose)
 
-        max_message_size_bytes = (await async_client_context.hello)["maxMessageSizeBytes"]
-
         # Document too large.
-        b_repeated = "b" * max_message_size_bytes
+        b_repeated = "b" * self.max_message_size_bytes
         models = [InsertOne(namespace="db.coll", document={"a": b_repeated})]
         with self.assertRaises(InvalidOperation) as context:
             await client.bulk_write(models=models)
         self.assertIn("cannot do an empty bulk write", context.exception._message)
 
         # Namespace too large.
-        c_repeated = "c" * max_message_size_bytes
+        c_repeated = "c" * self.max_message_size_bytes
         namespace = f"db.{c_repeated}"
         models = [InsertOne(namespace=namespace, document={"a": "b"})]
         with self.assertRaises(InvalidOperation) as context:
@@ -522,27 +524,32 @@ async def test_returns_error_if_auto_encryption_configured(self):
 
 # https://github.com/mongodb/specifications/blob/master/source/client-side-operations-timeout/tests/README.md#11-multi-batch-bulkwrites
 class TestClientBulkWriteTimeout(AsyncIntegrationTest):
+    async def asyncSetUp(self):
+        self.max_write_batch_size = await async_client_context.max_write_batch_size
+        self.max_bson_object_size = await async_client_context.max_bson_size
+        self.max_message_size_bytes = await async_client_context.max_message_size_bytes
+
     @async_client_context.require_version_min(8, 0, 0, -24)
     @async_client_context.require_failCommand_fail_point
     async def test_timeout_in_multi_batch_bulk_write(self):
+        _OVERHEAD = 500
+
         internal_client = await async_rs_or_single_client(timeoutMS=None)
         self.addAsyncCleanup(internal_client.aclose)
 
         collection = internal_client.db["coll"]
         self.addAsyncCleanup(collection.drop)
         await collection.drop()
 
-        max_bson_object_size = (await async_client_context.hello)["maxBsonObjectSize"]
-        max_message_size_bytes = (await async_client_context.hello)["maxMessageSizeBytes"]
         fail_command = {
             "configureFailPoint": "failCommand",
             "mode": {"times": 2},
             "data": {"failCommands": ["bulkWrite"], "blockConnection": True, "blockTimeMS": 1010},
         }
         async with self.fail_point(fail_command):
             models = []
-            num_models = int(max_message_size_bytes / max_bson_object_size + 1)
-            b_repeated = "b" * (max_bson_object_size - 500)
+            num_models = int(self.max_message_size_bytes / self.max_bson_object_size + 1)
+            b_repeated = "b" * (self.max_bson_object_size - _OVERHEAD)
             for _ in range(num_models):
                 models.append(
                     InsertOne(