Merge branch 'async-improvements' into PYTHON-4860

Author: NoahStapp

test/__init__.py

@@ -1115,26 +1115,10 @@ def enable_replication(self, client):
 class UnitTest(PyMongoTestCase):
     """Async base class for TestCases that don't require a connection to MongoDB."""
 
-    @classmethod
-    def setUpClass(cls):
-        if _IS_SYNC:
-            cls._setup_class()
-        else:
-            asyncio.run(cls._setup_class())
-
-    @classmethod
-    def tearDownClass(cls):
-        if _IS_SYNC:
-            cls._tearDown_class()
-        else:
-            asyncio.run(cls._tearDown_class())
-
-    @classmethod
-    def _setup_class(cls):
+    def setUp(self) -> None:
         pass
 
-    @classmethod
-    def _tearDown_class(cls):
+    def tearDown(self) -> None:
         pass
 
 
@@ -1145,37 +1129,20 @@ class IntegrationTest(PyMongoTestCase):
     db: Database
     credentials: Dict[str, str]
 
-    @classmethod
-    def setUpClass(cls):
-        if _IS_SYNC:
-            cls._setup_class()
-        else:
-            asyncio.run(cls._setup_class())
-
-    @classmethod
-    def tearDownClass(cls):
-        if _IS_SYNC:
-            cls._tearDown_class()
-        else:
-            asyncio.run(cls._tearDown_class())
-
-    @classmethod
     @client_context.require_connection
-    def _setup_class(cls):
-        if client_context.load_balancer and not getattr(cls, "RUN_ON_LOAD_BALANCER", False):
+    def setUp(self) -> None:
+        if not _IS_SYNC:
+            reset_client_context()
+        if client_context.load_balancer and not getattr(self, "RUN_ON_LOAD_BALANCER", False):
             raise SkipTest("this test does not support load balancers")
-        if client_context.serverless and not getattr(cls, "RUN_ON_SERVERLESS", False):
+        if client_context.serverless and not getattr(self, "RUN_ON_SERVERLESS", False):
             raise SkipTest("this test does not support serverless")
-        cls.client = client_context.client
-        cls.db = cls.client.pymongo_test
+        self.client = client_context.client
+        self.db = self.client.pymongo_test
         if client_context.auth_enabled:
-            cls.credentials = {"username": db_user, "password": db_pwd}
+            self.credentials = {"username": db_user, "password": db_pwd}
         else:
-            cls.credentials = {}
-
-    @classmethod
-    def _tearDown_class(cls):
-        pass
+            self.credentials = {}
 
     def cleanup_colls(self, *collections):
         """Cleanup collections faster than drop_collection."""
@@ -1201,37 +1168,14 @@ class MockClientTest(UnitTest):
     # MockClients tests that use replicaSet, directConnection=True, pass
     # multiple seed addresses, or wait for heartbeat events are incompatible
     # with loadBalanced=True.
-    @classmethod
-    def setUpClass(cls):
-        if _IS_SYNC:
-            cls._setup_class()
-        else:
-            asyncio.run(cls._setup_class())
-
-    @classmethod
-    def tearDownClass(cls):
-        if _IS_SYNC:
-            cls._tearDown_class()
-        else:
-            asyncio.run(cls._tearDown_class())
-
-    @classmethod
     @client_context.require_no_load_balancer
-    def _setup_class(cls):
-        pass
-
-    @classmethod
-    def _tearDown_class(cls):
-        pass
-
-    def setUp(self):
+    def setUp(self) -> None:
         super().setUp()
 
         self.client_knobs = client_knobs(heartbeat_frequency=0.001, min_heartbeat_interval=0.001)
-
         self.client_knobs.enable()
 
-    def tearDown(self):
+    def tearDown(self) -> None:
         self.client_knobs.disable()
         super().tearDown()
 

test/asynchronous/__init__.py

@@ -1133,26 +1133,10 @@ async def enable_replication(self, client):
 class AsyncUnitTest(AsyncPyMongoTestCase):
     """Async base class for TestCases that don't require a connection to MongoDB."""
 
-    @classmethod
-    def setUpClass(cls):
-        if _IS_SYNC:
-            cls._setup_class()
-        else:
-            asyncio.run(cls._setup_class())
-
-    @classmethod
-    def tearDownClass(cls):
-        if _IS_SYNC:
-            cls._tearDown_class()
-        else:
-            asyncio.run(cls._tearDown_class())
-
-    @classmethod
-    async def _setup_class(cls):
+    async def asyncSetUp(self) -> None:
         pass
 
-    @classmethod
-    async def _tearDown_class(cls):
+    async def asyncTearDown(self) -> None:
         pass
 
 
@@ -1163,37 +1147,20 @@ class AsyncIntegrationTest(AsyncPyMongoTestCase):
     db: AsyncDatabase
     credentials: Dict[str, str]
 
-    @classmethod
-    def setUpClass(cls):
-        if _IS_SYNC:
-            cls._setup_class()
-        else:
-            asyncio.run(cls._setup_class())
-
-    @classmethod
-    def tearDownClass(cls):
-        if _IS_SYNC:
-            cls._tearDown_class()
-        else:
-            asyncio.run(cls._tearDown_class())
-
-    @classmethod
     @async_client_context.require_connection
-    async def _setup_class(cls):
-        if async_client_context.load_balancer and not getattr(cls, "RUN_ON_LOAD_BALANCER", False):
+    async def asyncSetUp(self) -> None:
+        if not _IS_SYNC:
+            await reset_client_context()
+        if async_client_context.load_balancer and not getattr(self, "RUN_ON_LOAD_BALANCER", False):
             raise SkipTest("this test does not support load balancers")
-        if async_client_context.serverless and not getattr(cls, "RUN_ON_SERVERLESS", False):
+        if async_client_context.serverless and not getattr(self, "RUN_ON_SERVERLESS", False):
             raise SkipTest("this test does not support serverless")
-        cls.client = async_client_context.client
-        cls.db = cls.client.pymongo_test
+        self.client = async_client_context.client
+        self.db = self.client.pymongo_test
         if async_client_context.auth_enabled:
-            cls.credentials = {"username": db_user, "password": db_pwd}
+            self.credentials = {"username": db_user, "password": db_pwd}
         else:
-            cls.credentials = {}
-
-    @classmethod
-    async def _tearDown_class(cls):
-        pass
+            self.credentials = {}
 
     async def cleanup_colls(self, *collections):
         """Cleanup collections faster than drop_collection."""
@@ -1219,39 +1186,16 @@ class AsyncMockClientTest(AsyncUnitTest):
     # MockClients tests that use replicaSet, directConnection=True, pass
     # multiple seed addresses, or wait for heartbeat events are incompatible
     # with loadBalanced=True.
-    @classmethod
-    def setUpClass(cls):
-        if _IS_SYNC:
-            cls._setup_class()
-        else:
-            asyncio.run(cls._setup_class())
-
-    @classmethod
-    def tearDownClass(cls):
-        if _IS_SYNC:
-            cls._tearDown_class()
-        else:
-            asyncio.run(cls._tearDown_class())
-
-    @classmethod
     @async_client_context.require_no_load_balancer
-    async def _setup_class(cls):
-        pass
-
-    @classmethod
-    async def _tearDown_class(cls):
-        pass
-
-    def setUp(self):
-        super().setUp()
+    async def asyncSetUp(self) -> None:
+        await super().asyncSetUp()
 
         self.client_knobs = client_knobs(heartbeat_frequency=0.001, min_heartbeat_interval=0.001)
-
         self.client_knobs.enable()
 
-    def tearDown(self):
+    async def asyncTearDown(self) -> None:
         self.client_knobs.disable()
-        super().tearDown()
+        await super().asyncTearDown()
 
 
 async def async_setup():

test/asynchronous/test_bulk.py

@@ -42,15 +42,11 @@ class AsyncBulkTestBase(AsyncIntegrationTest):
     coll: AsyncCollection
     coll_w0: AsyncCollection
 
-    @classmethod
-    async def _setup_class(cls):
-        await super()._setup_class()
-        cls.coll = cls.db.test
-        cls.coll_w0 = cls.coll.with_options(write_concern=WriteConcern(w=0))
-
     async def asyncSetUp(self):
-        super().setUp()
+        await super().asyncSetUp()
+        self.coll = self.db.test
         await self.coll.drop()
+        self.coll_w0 = self.coll.with_options(write_concern=WriteConcern(w=0))
 
     def assertEqualResponse(self, expected, actual):
         """Compare response from bulk.execute() to expected response."""
@@ -787,14 +783,10 @@ async def test_large_inserts_unordered(self):
 
 
 class AsyncBulkAuthorizationTestBase(AsyncBulkTestBase):
-    @classmethod
     @async_client_context.require_auth
     @async_client_context.require_no_api_version
-    async def _setup_class(cls):
-        await super()._setup_class()
-
     async def asyncSetUp(self):
-        super().setUp()
+        await super().asyncSetUp()
         await async_client_context.create_user(self.db.name, "readonly", "pw", ["read"])
         await self.db.command(
             "createRole",
@@ -937,21 +929,19 @@ class AsyncTestBulkWriteConcern(AsyncBulkTestBase):
     w: Optional[int]
     secondary: AsyncMongoClient
 
-    @classmethod
-    async def _setup_class(cls):
-        await super()._setup_class()
-        cls.w = async_client_context.w
-        cls.secondary = None
-        if cls.w is not None and cls.w > 1:
+    async def asyncSetUp(self):
+        await super().asyncSetUp()
+        self.w = async_client_context.w
+        self.secondary = None
+        if self.w is not None and self.w > 1:
             for member in (await async_client_context.hello)["hosts"]:
                 if member != (await async_client_context.hello)["primary"]:
-                    cls.secondary = await cls.unmanaged_async_single_client(*partition_node(member))
+                    self.secondary = await self.async_single_client(*partition_node(member))
                     break
 
-    @classmethod
-    async def async_tearDownClass(cls):
-        if cls.secondary:
-            await cls.secondary.close()
+    async def asyncTearDown(self):
+        if self.secondary:
+            await self.secondary.close()
 
     async def cause_wtimeout(self, requests, ordered):
         if not async_client_context.test_commands_enabled:

test/asynchronous/test_change_stream.py

@@ -835,18 +835,16 @@ async def test_split_large_change(self):
 class TestClusterAsyncChangeStream(TestAsyncChangeStreamBase, APITestsMixin):
     dbs: list
 
-    @classmethod
     @async_client_context.require_version_min(4, 0, 0, -1)
     @async_client_context.require_change_streams
-    async def _setup_class(cls):
-        await super()._setup_class()
-        cls.dbs = [cls.db, cls.client.pymongo_test_2]
+    async def asyncSetUp(self) -> None:
+        await super().asyncSetUp()
+        self.dbs = [self.db, self.client.pymongo_test_2]
 
-    @classmethod
-    async def _tearDown_class(cls):
-        for db in cls.dbs:
-            await cls.client.drop_database(db)
-        await super()._tearDown_class()
+    async def asyncTearDown(self):
+        for db in self.dbs:
+            await self.client.drop_database(db)
+        await super().asyncTearDown()
 
     async def change_stream_with_client(self, client, *args, **kwargs):
         return await client.watch(*args, **kwargs)
@@ -897,11 +895,10 @@ async def test_full_pipeline(self):
 
 
 class TestAsyncDatabaseAsyncChangeStream(TestAsyncChangeStreamBase, APITestsMixin):
-    @classmethod
     @async_client_context.require_version_min(4, 0, 0, -1)
     @async_client_context.require_change_streams
-    async def _setup_class(cls):
-        await super()._setup_class()
+    async def asyncSetUp(self) -> None:
+        await super().asyncSetUp()
 
     async def change_stream_with_client(self, client, *args, **kwargs):
         return await client[self.db.name].watch(*args, **kwargs)
@@ -987,12 +984,9 @@ async def test_isolation(self):
 class TestAsyncCollectionAsyncChangeStream(
     TestAsyncChangeStreamBase, APITestsMixin, ProseSpecTestsMixin
 ):
-    @classmethod
     @async_client_context.require_change_streams
-    async def _setup_class(cls):
-        await super()._setup_class()
-
     async def asyncSetUp(self):
+        await super().asyncSetUp()
         # Use a new collection for each test.
         await self.watched_collection().drop()
         await self.watched_collection().insert_one({})
@@ -1132,20 +1126,11 @@ class TestAllLegacyScenarios(AsyncIntegrationTest):
     RUN_ON_LOAD_BALANCER = True
     listener: AllowListEventListener
 
-    @classmethod
     @async_client_context.require_connection
-    async def _setup_class(cls):
-        await super()._setup_class()
-        cls.listener = AllowListEventListener("aggregate", "getMore")
-        cls.client = await cls.unmanaged_async_rs_or_single_client(event_listeners=[cls.listener])
-
-    @classmethod
-    async def _tearDown_class(cls):
-        await cls.client.close()
-        await super()._tearDown_class()
-
-    def asyncSetUp(self):
-        super().asyncSetUp()
+    async def asyncSetUp(self):
+        await super().asyncSetUp()
+        self.listener = AllowListEventListener("aggregate", "getMore")
+        self.client = await self.async_rs_or_single_client(event_listeners=[self.listener])
         self.listener.reset()
 
     async def asyncSetUpCluster(self, scenario_dict):

test/asynchronous/test_client.py

@@ -130,16 +130,11 @@ class AsyncClientUnitTest(AsyncUnitTest):
 
     client: AsyncMongoClient
 
-    @classmethod
-    async def _setup_class(cls):
-        cls.client = await cls.unmanaged_async_rs_or_single_client(
+    async def asyncSetUp(self) -> None:
+        self.client = await self.async_rs_or_single_client(
             connect=False, serverSelectionTimeoutMS=100
         )
 
-    @classmethod
-    async def _tearDown_class(cls):
-        await cls.client.close()
-
     @pytest.fixture(autouse=True)
     def inject_fixtures(self, caplog):
         self._caplog = caplog