Merge branch 'master' into PYTHON-5867

Author: NoahStapp

pymongo/asynchronous/pool.py

@@ -1065,10 +1065,16 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
                     reason=_verbose_connection_error_reason(ConnectionClosedReason.ERROR),
                     error=ConnectionClosedReason.ERROR,
                 )
-            self._handle_connection_error(error)
             if isinstance(error, (IOError, OSError, *SSLErrors)):
                 details = _get_timeout_details(self.opts)
-                _raise_connection_failure(self.address, error, timeout_details=details)
+                # Wrap to AutoReconnect/NetworkTimeout BEFORE labeling so the
+                # SystemOverloadedError label lands on the propagated exception.
+                try:
+                    _raise_connection_failure(self.address, error, timeout_details=details)
+                except (AutoReconnect, NetworkTimeout) as wrapped:
+                    self._handle_connection_error(wrapped)
+                    raise
+            self._handle_connection_error(error)
             raise
 
         conn = AsyncConnection(networking_interface, self, self.address, conn_id, self.is_sdam)  # type: ignore[arg-type]

pymongo/synchronous/pool.py

@@ -1061,10 +1061,16 @@ def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> Connect
                     reason=_verbose_connection_error_reason(ConnectionClosedReason.ERROR),
                     error=ConnectionClosedReason.ERROR,
                 )
-            self._handle_connection_error(error)
             if isinstance(error, (IOError, OSError, *SSLErrors)):
                 details = _get_timeout_details(self.opts)
-                _raise_connection_failure(self.address, error, timeout_details=details)
+                # Wrap to AutoReconnect/NetworkTimeout BEFORE labeling so the
+                # SystemOverloadedError label lands on the propagated exception.
+                try:
+                    _raise_connection_failure(self.address, error, timeout_details=details)
+                except (AutoReconnect, NetworkTimeout) as wrapped:
+                    self._handle_connection_error(wrapped)
+                    raise
+            self._handle_connection_error(error)
             raise
 
         conn = Connection(networking_interface, self, self.address, conn_id, self.is_sdam)  # type: ignore[arg-type]

test/asynchronous/test_encryption.py

@@ -2207,6 +2207,13 @@ async def test_01_insert_encrypted_indexed_and_find(self):
         self.assertEqual(docs[0]["encryptedIndexed"], val)
 
     async def test_02_insert_encrypted_indexed_and_find_contention(self):
+        # setUp creates the collection with contention=0 (from encryptedFields.json).
+        # This test uses contention_factor=10, so recreate the collection with contention=10.
+        await self.db.drop_collection("explicit_encryption", encrypted_fields=self.encrypted_fields)
+        encrypted_fields = copy.deepcopy(self.encrypted_fields)
+        encrypted_fields["fields"][0]["queries"]["contention"] = 10
+        await self.db.command("create", "explicit_encryption", encryptedFields=encrypted_fields)
+
         val = "encrypted indexed value"
         contention = 10
         for _ in range(contention):
@@ -2217,20 +2224,7 @@ async def test_02_insert_encrypted_indexed_and_find_contention(self):
                 {"encryptedIndexed": insert_payload}
             )
 
-        find_payload = await self.client_encryption.encrypt(
-            val, Algorithm.INDEXED, self.key1_id, query_type=QueryType.EQUALITY, contention_factor=0
-        )
-        docs = (
-            await self.encrypted_client[self.db.name]
-            .explicit_encryption.find({"encryptedIndexed": find_payload})
-            .to_list()
-        )
-
-        self.assertLessEqual(len(docs), 10)
-        for doc in docs:
-            self.assertEqual(doc["encryptedIndexed"], val)
-
-        # Find with contention_factor will return all 10 documents.
+        # Find with matching contention_factor returns all 10 documents.
         find_payload = await self.client_encryption.encrypt(
             val,
             Algorithm.INDEXED,

test/test_encryption.py

@@ -2195,6 +2195,13 @@ def test_01_insert_encrypted_indexed_and_find(self):
         self.assertEqual(docs[0]["encryptedIndexed"], val)
 
     def test_02_insert_encrypted_indexed_and_find_contention(self):
+        # setUp creates the collection with contention=0 (from encryptedFields.json).
+        # This test uses contention_factor=10, so recreate the collection with contention=10.
+        self.db.drop_collection("explicit_encryption", encrypted_fields=self.encrypted_fields)
+        encrypted_fields = copy.deepcopy(self.encrypted_fields)
+        encrypted_fields["fields"][0]["queries"]["contention"] = 10
+        self.db.command("create", "explicit_encryption", encryptedFields=encrypted_fields)
+
         val = "encrypted indexed value"
         contention = 10
         for _ in range(contention):
@@ -2205,20 +2212,7 @@ def test_02_insert_encrypted_indexed_and_find_contention(self):
                 {"encryptedIndexed": insert_payload}
             )
 
-        find_payload = self.client_encryption.encrypt(
-            val, Algorithm.INDEXED, self.key1_id, query_type=QueryType.EQUALITY, contention_factor=0
-        )
-        docs = (
-            self.encrypted_client[self.db.name]
-            .explicit_encryption.find({"encryptedIndexed": find_payload})
-            .to_list()
-        )
-
-        self.assertLessEqual(len(docs), 10)
-        for doc in docs:
-            self.assertEqual(doc["encryptedIndexed"], val)
-
-        # Find with contention_factor will return all 10 documents.
+        # Find with matching contention_factor returns all 10 documents.
         find_payload = self.client_encryption.encrypt(
             val,
             Algorithm.INDEXED,