Use one loop for asyncio test suite

Author: ShaneHarvey

test/__init__.py

@@ -30,28 +30,6 @@
 import unittest
 import warnings
 from asyncio import iscoroutinefunction
-from test.helpers import (
-    COMPRESSORS,
-    IS_SRV,
-    MONGODB_API_VERSION,
-    MULTI_MONGOS_LB_URI,
-    TEST_LOADBALANCER,
-    TEST_SERVERLESS,
-    TLS_OPTIONS,
-    SystemCertsPatcher,
-    client_knobs,
-    db_pwd,
-    db_user,
-    global_knobs,
-    host,
-    is_server_resolvable,
-    port,
-    print_running_topology,
-    print_thread_stacks,
-    print_thread_tracebacks,
-    sanitize_cmd,
-    sanitize_reply,
-)
 
 from pymongo.uri_parser import parse_uri
 
@@ -63,7 +41,6 @@
     HAVE_IPADDRESS = False
 from contextlib import contextmanager
 from functools import partial, wraps
-from test.version import Version
 from typing import Any, Callable, Dict, Generator, overload
 from unittest import SkipTest
 from urllib.parse import quote_plus
@@ -78,6 +55,32 @@
 from pymongo.synchronous.database import Database
 from pymongo.synchronous.mongo_client import MongoClient
 
+sys.path[0:0] = [""]
+
+from test.helpers import (
+    COMPRESSORS,
+    IS_SRV,
+    MONGODB_API_VERSION,
+    MULTI_MONGOS_LB_URI,
+    TEST_LOADBALANCER,
+    TEST_SERVERLESS,
+    TLS_OPTIONS,
+    SystemCertsPatcher,
+    client_knobs,
+    db_pwd,
+    db_user,
+    global_knobs,
+    host,
+    is_server_resolvable,
+    port,
+    print_running_topology,
+    print_thread_stacks,
+    print_thread_tracebacks,
+    sanitize_cmd,
+    sanitize_reply,
+)
+from test.version import Version
+
 _IS_SYNC = True
 
 
@@ -875,6 +878,40 @@ def reset_client_context():
 
 
 class PyMongoTestCase(unittest.TestCase):
+    if not _IS_SYNC:
+        # Customize async TestCase to use a single event loop for all tests.
+        def __init__(self, methodName="runTest"):
+            super().__init__(methodName)
+            try:
+                self.loop = asyncio.get_event_loop()
+            except RuntimeError:
+                self.loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(self.loop)
+
+        def setUp(self):
+            pass
+
+        def tearDown(self):
+            pass
+
+        # See TestCase.addCleanup.
+        def addCleanup(self, func, /, *args, **kwargs):
+            self.addCleanup(*(func, *args), **kwargs)
+
+        def run(self, result=None):
+            if result is None:
+                result = self.defaultTestResult()
+            result.startTest(self)
+
+            try:
+                self.setUp()
+                self.loop.run_until_complete(self.setUp())
+                self.loop.run_until_complete(getattr(self, self._testMethodName)())
+                self.loop.run_until_complete(self.tearDown())
+            finally:
+                self.tearDown()
+                result.stopTest(self)
+
     def assertEqualCommand(self, expected, actual, msg=None):
         self.assertEqual(sanitize_cmd(expected), sanitize_cmd(actual), msg)
 

test/asynchronous/__init__.py

@@ -30,28 +30,6 @@
 import unittest
 import warnings
 from asyncio import iscoroutinefunction
-from test.helpers import (
-    COMPRESSORS,
-    IS_SRV,
-    MONGODB_API_VERSION,
-    MULTI_MONGOS_LB_URI,
-    TEST_LOADBALANCER,
-    TEST_SERVERLESS,
-    TLS_OPTIONS,
-    SystemCertsPatcher,
-    client_knobs,
-    db_pwd,
-    db_user,
-    global_knobs,
-    host,
-    is_server_resolvable,
-    port,
-    print_running_topology,
-    print_thread_stacks,
-    print_thread_tracebacks,
-    sanitize_cmd,
-    sanitize_reply,
-)
 
 from pymongo.uri_parser import parse_uri
 
@@ -63,7 +41,6 @@
     HAVE_IPADDRESS = False
 from contextlib import asynccontextmanager, contextmanager
 from functools import partial, wraps
-from test.version import Version
 from typing import Any, Callable, Dict, Generator, overload
 from unittest import SkipTest
 from urllib.parse import quote_plus
@@ -78,6 +55,32 @@
 from pymongo.server_api import ServerApi
 from pymongo.ssl_support import HAVE_SSL, _ssl  # type:ignore[attr-defined]
 
+sys.path[0:0] = [""]
+
+from test.helpers import (
+    COMPRESSORS,
+    IS_SRV,
+    MONGODB_API_VERSION,
+    MULTI_MONGOS_LB_URI,
+    TEST_LOADBALANCER,
+    TEST_SERVERLESS,
+    TLS_OPTIONS,
+    SystemCertsPatcher,
+    client_knobs,
+    db_pwd,
+    db_user,
+    global_knobs,
+    host,
+    is_server_resolvable,
+    port,
+    print_running_topology,
+    print_thread_stacks,
+    print_thread_tracebacks,
+    sanitize_cmd,
+    sanitize_reply,
+)
+from test.version import Version
+
 _IS_SYNC = False
 
 
@@ -876,7 +879,41 @@ async def reset_client_context():
     await async_client_context._init_client()
 
 
-class AsyncPyMongoTestCase(unittest.IsolatedAsyncioTestCase):
+class AsyncPyMongoTestCase(unittest.TestCase):
+    if not _IS_SYNC:
+        # Customize async TestCase to use a single event loop for all tests.
+        def __init__(self, methodName="runTest"):
+            super().__init__(methodName)
+            try:
+                self.loop = asyncio.get_event_loop()
+            except RuntimeError:
+                self.loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(self.loop)
+
+        async def asyncSetUp(self):
+            pass
+
+        async def asyncTearDown(self):
+            pass
+
+        # See IsolatedAsyncioTestCase.addAsyncCleanup.
+        def addAsyncCleanup(self, func, /, *args, **kwargs):
+            self.addCleanup(*(func, *args), **kwargs)
+
+        def run(self, result=None):
+            if result is None:
+                result = self.defaultTestResult()
+            result.startTest(self)
+
+            try:
+                self.setUp()
+                self.loop.run_until_complete(self.asyncSetUp())
+                self.loop.run_until_complete(getattr(self, self._testMethodName)())
+                self.loop.run_until_complete(self.asyncTearDown())
+            finally:
+                self.tearDown()
+                result.stopTest(self)
+
     def assertEqualCommand(self, expected, actual, msg=None):
         self.assertEqual(sanitize_cmd(expected), sanitize_cmd(actual), msg)