Add test

Author: NoahStapp

test/__init__.py

@@ -826,6 +826,14 @@ def require_sync(self, func):
             lambda: _IS_SYNC, "This test only works with the synchronous API", func=func
         )
 
+    def require_async(self, func):
+        """Run a test only if using the asynchronous API."""  # unasync: off
+        return self._require(
+            lambda: not _IS_SYNC,
+            "This test only works with the asynchronous API",  # unasync: off
+            func=func,
+        )
+
     def mongos_seeds(self):
         return ",".join("{}:{}".format(*address) for address in self.mongoses)
 

test/asynchronous/__init__.py

@@ -828,6 +828,14 @@ def require_sync(self, func):
             lambda: _IS_SYNC, "This test only works with the synchronous API", func=func
         )
 
+    def require_async(self, func):
+        """Run a test only if using the asynchronous API."""  # unasync: off
+        return self._require(
+            lambda: not _IS_SYNC,
+            "This test only works with the asynchronous API",  # unasync: off
+            func=func,
+        )
+
     def mongos_seeds(self):
         return ",".join("{}:{}".format(*address) for address in self.mongoses)
 

test/asynchronous/test_discovery_and_monitoring.py

@@ -20,10 +20,15 @@
 import socketserver
 import sys
 import threading
+import time
 from asyncio import StreamReader, StreamWriter
 from pathlib import Path
 from test.asynchronous.helpers import ConcurrentRunner
 
+from pymongo.asynchronous.pool import AsyncConnection
+from pymongo.operations import _Op
+from pymongo.server_selectors import readable_server_selector
+
 sys.path[0:0] = [""]
 
 from test.asynchronous import (
@@ -46,6 +51,7 @@
     async_barrier_wait,
     async_create_barrier,
     async_wait_until,
+    delay,
     server_name_to_type,
 )
 from unittest.mock import patch
@@ -370,6 +376,69 @@ async def test_pool_unpause(self):
             await listener.async_wait_for_event(monitoring.ServerHeartbeatSucceededEvent, 1)
             await listener.async_wait_for_event(monitoring.PoolReadyEvent, 1)
 
+    @async_client_context.require_failCommand_appName
+    @async_client_context.require_test_commands
+    @async_client_context.require_async
+    async def test_connection_close_does_not_block_other_operations(self):
+        listener = CMAPHeartbeatListener()
+        client = await self.async_single_client(
+            appName="SDAMConnectionCloseTest",
+            event_listeners=[listener],
+            heartbeatFrequencyMS=500,
+            minPoolSize=10,
+        )
+        server = await (await client._get_topology()).select_server(
+            readable_server_selector, _Op.TEST
+        )
+        await async_wait_until(
+            lambda: len(server._pool.conns) == 10,
+            "pool initialized with 10 connections",
+        )
+
+        await client.db.test.insert_one({"x": 1})
+        close_delay = 0.05
+        latencies = []
+
+        async def run_task():
+            while True:
+                start_time = time.monotonic()
+                await client.db.test.find_one({})
+                elapsed = time.monotonic() - start_time
+                latencies.append(elapsed)
+                if elapsed >= close_delay:
+                    break
+                await asyncio.sleep(0.001)
+
+        task = ConcurrentRunner(target=run_task)
+        await task.start()
+        original_close = AsyncConnection.close_conn
+        try:
+            # Artificially delay the close operation to simulate a slow close
+            async def mock_close(self, reason):
+                await asyncio.sleep(close_delay)
+                await original_close(self, reason)
+
+            AsyncConnection.close_conn = mock_close
+
+            fail_hello = {
+                "mode": {"times": 4},
+                "data": {
+                    "failCommands": [HelloCompat.LEGACY_CMD, "hello"],
+                    "errorCode": 91,
+                    "appName": "SDAMConnectionCloseTest",
+                },
+            }
+            async with self.fail_point(fail_hello):
+                # Wait for server heartbeat to fail
+                await listener.async_wait_for_event(monitoring.ServerHeartbeatFailedEvent, 1)
+            # Wait until all idle connections are closed to simulate real-world conditions
+            await listener.async_wait_for_event(monitoring.ConnectionClosedEvent, 10)
+            # No operation latency should not significantly exceed close_delay
+            self.assertLessEqual(max(latencies), close_delay * 1.5)
+        finally:
+            AsyncConnection.close_conn = original_close
+            await task.join()
+
 
 class TestServerMonitoringMode(AsyncIntegrationTest):
     @async_client_context.require_no_serverless

test/test_discovery_and_monitoring.py

@@ -20,10 +20,15 @@
 import socketserver
 import sys
 import threading
+import time
 from asyncio import StreamReader, StreamWriter
 from pathlib import Path
 from test.helpers import ConcurrentRunner
 
+from pymongo.operations import _Op
+from pymongo.server_selectors import readable_server_selector
+from pymongo.synchronous.pool import Connection
+
 sys.path[0:0] = [""]
 
 from test import (
@@ -45,6 +50,7 @@
     assertion_context,
     barrier_wait,
     create_barrier,
+    delay,
     server_name_to_type,
     wait_until,
 )
@@ -370,6 +376,67 @@ def test_pool_unpause(self):
             listener.wait_for_event(monitoring.ServerHeartbeatSucceededEvent, 1)
             listener.wait_for_event(monitoring.PoolReadyEvent, 1)
 
+    @client_context.require_failCommand_appName
+    @client_context.require_test_commands
+    @client_context.require_async
+    def test_connection_close_does_not_block_other_operations(self):
+        listener = CMAPHeartbeatListener()
+        client = self.single_client(
+            appName="SDAMConnectionCloseTest",
+            event_listeners=[listener],
+            heartbeatFrequencyMS=500,
+            minPoolSize=10,
+        )
+        server = (client._get_topology()).select_server(readable_server_selector, _Op.TEST)
+        wait_until(
+            lambda: len(server._pool.conns) == 10,
+            "pool initialized with 10 connections",
+        )
+
+        client.db.test.insert_one({"x": 1})
+        close_delay = 0.05
+        latencies = []
+
+        def run_task():
+            while True:
+                start_time = time.monotonic()
+                client.db.test.find_one({})
+                elapsed = time.monotonic() - start_time
+                latencies.append(elapsed)
+                if elapsed >= close_delay:
+                    break
+                time.sleep(0.001)
+
+        task = ConcurrentRunner(target=run_task)
+        task.start()
+        original_close = Connection.close_conn
+        try:
+            # Artificially delay the close operation to simulate a slow close
+            def mock_close(self, reason):
+                time.sleep(close_delay)
+                original_close(self, reason)
+
+            Connection.close_conn = mock_close
+
+            fail_hello = {
+                "mode": {"times": 4},
+                "data": {
+                    "failCommands": [HelloCompat.LEGACY_CMD, "hello"],
+                    "errorCode": 91,
+                    "appName": "SDAMConnectionCloseTest",
+                },
+            }
+            with self.fail_point(fail_hello):
+                # Wait for server heartbeat to fail
+                listener.wait_for_event(monitoring.ServerHeartbeatFailedEvent, 1)
+            # Wait until all idle connections are closed to simulate real-world conditions
+            listener.wait_for_event(monitoring.ConnectionClosedEvent, 10)
+            # No operation latency should not significantly exceed close_delay
+            self.assertLessEqual(max(latencies), close_delay * 1.5)
+        finally:
+            Connection.close_conn = original_close
+            task.join()
+
 
 class TestServerMonitoringMode(IntegrationTest):
     @client_context.require_no_serverless

tools/synchro.py

@@ -288,7 +288,8 @@ def process_files(
                 if file in docstring_translate_files:
                     lines = translate_docstrings(lines)
                 if file in sync_test_files:
-                    translate_imports(lines)
+                    lines = translate_imports(lines)
+                lines = process_ignores(lines)
                 f.seek(0)
                 f.writelines(lines)
                 f.truncate()
@@ -390,6 +391,14 @@ def translate_docstrings(lines: list[str]) -> list[str]:
     return [line for line in lines if line != "DOCSTRING_REMOVED"]
 
 
+def process_ignores(lines: list[str]) -> list[str]:
+    for i in range(len(lines)):
+        for k, v in replacements.items():
+            if "unasync: off" in lines[i] and v in lines[i]:
+                lines[i] = lines[i].replace(v, k)
+    return lines
+
+
 def unasync_directory(files: list[str], src: str, dest: str, replacements: dict[str, str]) -> None:
     unasync_files(
         files,