Test cleanup

Author: NoahStapp

pymongo/asynchronous/pool.py

@@ -445,7 +445,7 @@ async def command(
             raise
         # Catch socket.error, KeyboardInterrupt, CancelledError, etc. and close ourselves.
         except BaseException as error:
-            await _raise_connection_failure(error)
+            await self._raise_connection_failure(error)
 
     async def send_message(self, message: bytes, max_doc_size: int) -> None:
         """Send a raw BSON message or raise ConnectionFailure.
@@ -462,7 +462,7 @@ async def send_message(self, message: bytes, max_doc_size: int) -> None:
             await async_sendall(self.conn.get_conn, message)
         # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException as error:
-            await _raise_connection_failure(error)
+            await self._raise_connection_failure(error)
 
     async def receive_message(self, request_id: Optional[int]) -> Union[_OpReply, _OpMsg]:
         """Receive a raw BSON message or raise ConnectionFailure.
@@ -473,7 +473,7 @@ async def receive_message(self, request_id: Optional[int]) -> Union[_OpReply, _O
             return await async_receive_message(self, request_id, self.max_message_size)
         # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException as error:
-            await _raise_connection_failure(error)
+            await self._raise_connection_failure(error)
 
     def _raise_if_not_writable(self, unacknowledged: bool) -> None:
         """Raise NotPrimaryError on unacknowledged write if this socket is not

pymongo/synchronous/pool.py

@@ -445,7 +445,7 @@ def command(
             raise
         # Catch socket.error, KeyboardInterrupt, CancelledError, etc. and close ourselves.
         except BaseException as error:
-            _raise_connection_failure(error)
+            self._raise_connection_failure(error)
 
     def send_message(self, message: bytes, max_doc_size: int) -> None:
         """Send a raw BSON message or raise ConnectionFailure.
@@ -462,7 +462,7 @@ def send_message(self, message: bytes, max_doc_size: int) -> None:
             sendall(self.conn.get_conn, message)
         # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException as error:
-            _raise_connection_failure(error)
+            self._raise_connection_failure(error)
 
     def receive_message(self, request_id: Optional[int]) -> Union[_OpReply, _OpMsg]:
         """Receive a raw BSON message or raise ConnectionFailure.
@@ -473,7 +473,7 @@ def receive_message(self, request_id: Optional[int]) -> Union[_OpReply, _OpMsg]:
             return receive_message(self, request_id, self.max_message_size)
         # Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
         except BaseException as error:
-            _raise_connection_failure(error)
+            self._raise_connection_failure(error)
 
     def _raise_if_not_writable(self, unacknowledged: bool) -> None:
         """Raise NotPrimaryError on unacknowledged write if this socket is not

test/asynchronous/test_bulk.py

@@ -505,7 +505,7 @@ async def test_single_ordered_batch(self):
 
     async def test_single_error_ordered_batch(self):
         await self.coll.create_index("a", unique=True)
-        self.addToCleanup(self.coll.drop_index, [("a", 1)])
+        self.addAsyncCleanup(self.coll.drop_index, [("a", 1)])
         requests: list = [
             InsertOne({"b": 1, "a": 1}),
             UpdateOne({"b": 2}, {"$set": {"a": 1}}, upsert=True),
@@ -547,7 +547,7 @@ async def test_single_error_ordered_batch(self):
 
     async def test_multiple_error_ordered_batch(self):
         await self.coll.create_index("a", unique=True)
-        self.addToCleanup(self.coll.drop_index, [("a", 1)])
+        self.addAsyncCleanup(self.coll.drop_index, [("a", 1)])
         requests: list = [
             InsertOne({"b": 1, "a": 1}),
             UpdateOne({"b": 2}, {"$set": {"a": 1}}, upsert=True),
@@ -616,7 +616,7 @@ async def test_single_unordered_batch(self):
 
     async def test_single_error_unordered_batch(self):
         await self.coll.create_index("a", unique=True)
-        self.addToCleanup(self.coll.drop_index, [("a", 1)])
+        self.addAsyncCleanup(self.coll.drop_index, [("a", 1)])
         requests: list = [
             InsertOne({"b": 1, "a": 1}),
             UpdateOne({"b": 2}, {"$set": {"a": 1}}, upsert=True),
@@ -659,7 +659,7 @@ async def test_single_error_unordered_batch(self):
 
     async def test_multiple_error_unordered_batch(self):
         await self.coll.create_index("a", unique=True)
-        self.addToCleanup(self.coll.drop_index, [("a", 1)])
+        self.addAsyncCleanup(self.coll.drop_index, [("a", 1)])
         requests: list = [
             InsertOne({"b": 1, "a": 1}),
             UpdateOne({"b": 2}, {"$set": {"a": 3}}, upsert=True),
@@ -1002,7 +1002,7 @@ async def test_write_concern_failure_ordered(self):
 
         await self.coll.delete_many({})
         await self.coll.create_index("a", unique=True)
-        self.addToCleanup(self.coll.drop_index, [("a", 1)])
+        self.addAsyncCleanup(self.coll.drop_index, [("a", 1)])
 
         # Fail due to write concern support as well
         # as duplicate key error on ordered batch.
@@ -1077,7 +1077,7 @@ async def test_write_concern_failure_unordered(self):
 
         await self.coll.delete_many({})
         await self.coll.create_index("a", unique=True)
-        self.addToCleanup(self.coll.drop_index, [("a", 1)])
+        self.addAsyncCleanup(self.coll.drop_index, [("a", 1)])
 
         # Fail due to write concern support as well
         # as duplicate key error on unordered batch.

test/asynchronous/test_change_stream.py

@@ -165,7 +165,7 @@ async def test_try_next(self):
         coll = self.watched_collection().with_options(write_concern=WriteConcern("majority"))
         await coll.drop()
         await coll.insert_one({})
-        self.addToCleanup(coll.drop)
+        self.addAsyncCleanup(coll.drop)
         async with await self.change_stream(max_await_time_ms=250) as stream:
             self.assertIsNone(await stream.try_next())  # No changes initially.
             await coll.insert_one({})  # Generate a change.
@@ -191,7 +191,7 @@ async def test_try_next_runs_one_getmore(self):
         # Create the watched collection before starting the change stream to
         # skip any "create" events.
         await coll.insert_one({"_id": 1})
-        self.addToCleanup(coll.drop)
+        self.addAsyncCleanup(coll.drop)
         async with await self.change_stream_with_client(client, max_await_time_ms=250) as stream:
             self.assertEqual(listener.started_command_names(), ["aggregate"])
             listener.reset()
@@ -249,7 +249,7 @@ async def test_batch_size_is_honored(self):
         # Create the watched collection before starting the change stream to
         # skip any "create" events.
         await coll.insert_one({"_id": 1})
-        self.addToCleanup(coll.drop)
+        self.addAsyncCleanup(coll.drop)
         # Expected batchSize.
         expected = {"batchSize": 23}
         async with await self.change_stream_with_client(
@@ -489,7 +489,7 @@ async def _client_with_listener(self, *commands):
         client = await AsyncPyMongoTestCase.unmanaged_async_rs_or_single_client(
             event_listeners=[listener]
         )
-        self.addToCleanup(client.close)
+        self.addAsyncCleanup(client.close)
         return client, listener
 
     @no_type_check
@@ -1156,7 +1156,7 @@ async def setFailPoint(self, scenario_dict):
         fail_cmd = SON([("configureFailPoint", "failCommand")])
         fail_cmd.update(fail_point)
         await async_client_context.client.admin.command(fail_cmd)
-        self.addToCleanup(
+        self.addAsyncCleanup(
             async_client_context.client.admin.command,
             "configureFailPoint",
             fail_cmd["configureFailPoint"],

test/asynchronous/test_client.py

@@ -1102,8 +1102,8 @@ def test_bad_uri(self):
     async def test_auth_from_uri(self):
         host, port = await async_client_context.host, await async_client_context.port
         await async_client_context.create_user("admin", "admin", "pass")
-        self.addToCleanup(async_client_context.drop_user, "admin", "admin")
-        self.addToCleanup(remove_all_users, self.client.pymongo_test)
+        self.addAsyncCleanup(async_client_context.drop_user, "admin", "admin")
+        self.addAsyncCleanup(remove_all_users, self.client.pymongo_test)
 
         await async_client_context.create_user(
             "pymongo_test", "user", "pass", roles=["userAdmin", "readWrite"]
@@ -1149,7 +1149,7 @@ async def test_auth_from_uri(self):
     @async_client_context.require_auth
     async def test_username_and_password(self):
         await async_client_context.create_user("admin", "ad min", "pa/ss")
-        self.addToCleanup(async_client_context.drop_user, "admin", "ad min")
+        self.addAsyncCleanup(async_client_context.drop_user, "admin", "ad min")
 
         c = await self.async_rs_or_single_client_noauth(username="ad min", password="pa/ss")
 
@@ -1258,7 +1258,6 @@ async def test_socket_timeout(self):
         no_timeout = self.client
         timeout_sec = 1
         timeout = await self.async_rs_or_single_client(socketTimeoutMS=1000 * timeout_sec)
-        self.addToCleanup(timeout.close)
 
         await no_timeout.pymongo_test.drop_collection("test")
         await no_timeout.pymongo_test.test.insert_one({"x": 1})
@@ -1325,7 +1324,7 @@ async def test_tz_aware(self):
         self.assertRaises(ValueError, AsyncMongoClient, tz_aware="foo")
 
         aware = await self.async_rs_or_single_client(tz_aware=True)
-        self.addToCleanup(aware.close)
+        self.addAsyncCleanup(aware.close)
         naive = self.client
         await aware.pymongo_test.drop_collection("test")
 
@@ -1477,7 +1476,7 @@ async def test_lazy_connect_w0(self):
         # Use a separate collection to avoid races where we're still
         # completing an operation on a collection while the next test begins.
         await async_client_context.client.drop_database("test_lazy_connect_w0")
-        self.addToCleanup(async_client_context.client.drop_database, "test_lazy_connect_w0")
+        self.addAsyncCleanup(async_client_context.client.drop_database, "test_lazy_connect_w0")
 
         client = await self.async_rs_or_single_client(connect=False, w=0)
         await client.test_lazy_connect_w0.test.insert_one({})
@@ -2160,7 +2159,7 @@ async def test_exhaust_getmore_server_error(self):
         await collection.drop()
 
         await collection.insert_many([{} for _ in range(200)])
-        self.addToCleanup(async_client_context.client.pymongo_test.test.drop)
+        self.addAsyncCleanup(async_client_context.client.pymongo_test.test.drop)
 
         pool = await async_get_pool(client)
         pool._check_interval_seconds = None  # Never check.
@@ -2407,7 +2406,7 @@ async def test_discover_primary(self):
             replicaSet="rs",
             heartbeatFrequencyMS=500,
         )
-        self.addToCleanup(c.close)
+        self.addAsyncCleanup(c.close)
 
         await async_wait_until(lambda: len(c.nodes) == 3, "connect")
 
@@ -2434,7 +2433,7 @@ async def test_reconnect(self):
             retryReads=False,
             serverSelectionTimeoutMS=1000,
         )
-        self.addToCleanup(c.close)
+        self.addAsyncCleanup(c.close)
 
         await async_wait_until(lambda: len(c.nodes) == 3, "connect")
 
@@ -2472,7 +2471,7 @@ async def _test_network_error(self, operation_callback):
                 serverSelectionTimeoutMS=1000,
             )
 
-            self.addToCleanup(c.close)
+            self.addAsyncCleanup(c.close)
 
             # Set host-specific information so we can test whether it is reset.
             c.set_wire_version_range("a:1", 2, MIN_SUPPORTED_WIRE_VERSION)
@@ -2548,7 +2547,7 @@ async def test_rs_client_does_not_maintain_pool_to_arbiters(self):
             minPoolSize=1,  # minPoolSize
             event_listeners=[listener],
         )
-        self.addToCleanup(c.close)
+        self.addAsyncCleanup(c.close)
 
         await async_wait_until(lambda: len(c.nodes) == 3, "connect")
         self.assertEqual(await c.address, ("a", 1))
@@ -2578,7 +2577,7 @@ async def test_direct_client_maintains_pool_to_arbiter(self):
             minPoolSize=1,  # minPoolSize
             event_listeners=[listener],
         )
-        self.addToCleanup(c.close)
+        self.addAsyncCleanup(c.close)
 
         await async_wait_until(lambda: len(c.nodes) == 1, "connect")
         self.assertEqual(await c.address, ("c", 3))

test/asynchronous/test_client_bulk_write.py

@@ -102,7 +102,7 @@ async def test_batch_splits_if_num_operations_too_large(self):
         models = []
         for _ in range(self.max_write_batch_size + 1):
             models.append(InsertOne(namespace="db.coll", document={"a": "b"}))
-        self.addToCleanup(client.db["coll"].drop)
+        self.addAsyncCleanup(client.db["coll"].drop)
 
         result = await client.bulk_write(models=models)
         self.assertEqual(result.inserted_count, self.max_write_batch_size + 1)
@@ -134,7 +134,7 @@ async def test_batch_splits_if_ops_payload_too_large(self):
                     document={"a": b_repeated},
                 )
             )
-        self.addToCleanup(client.db["coll"].drop)
+        self.addAsyncCleanup(client.db["coll"].drop)
 
         result = await client.bulk_write(models=models)
         self.assertEqual(result.inserted_count, num_models)
@@ -177,7 +177,7 @@ async def test_collects_write_concern_errors_across_batches(self):
                         document={"a": "b"},
                     )
                 )
-            self.addToCleanup(client.db["coll"].drop)
+            self.addAsyncCleanup(client.db["coll"].drop)
 
             with self.assertRaises(ClientBulkWriteException) as context:
                 await client.bulk_write(models=models)
@@ -200,7 +200,7 @@ async def test_collects_write_errors_across_batches_unordered(self):
         client = await self.async_rs_or_single_client(event_listeners=[listener])
 
         collection = client.db["coll"]
-        self.addToCleanup(collection.drop)
+        self.addAsyncCleanup(collection.drop)
         await collection.drop()
         await collection.insert_one(document={"_id": 1})
 
@@ -230,7 +230,7 @@ async def test_collects_write_errors_across_batches_ordered(self):
         client = await self.async_rs_or_single_client(event_listeners=[listener])
 
         collection = client.db["coll"]
-        self.addToCleanup(collection.drop)
+        self.addAsyncCleanup(collection.drop)
         await collection.drop()
         await collection.insert_one(document={"_id": 1})
 
@@ -260,7 +260,7 @@ async def test_handles_cursor_requiring_getMore(self):
         client = await self.async_rs_or_single_client(event_listeners=[listener])
 
         collection = client.db["coll"]
-        self.addToCleanup(collection.drop)
+        self.addAsyncCleanup(collection.drop)
         await collection.drop()
 
         models = []
@@ -301,7 +301,7 @@ async def test_handles_cursor_requiring_getMore_within_transaction(self):
         client = await self.async_rs_or_single_client(event_listeners=[listener])
 
         collection = client.db["coll"]
-        self.addToCleanup(collection.drop)
+        self.addAsyncCleanup(collection.drop)
         await collection.drop()
 
         async with client.start_session() as session:
@@ -344,7 +344,7 @@ async def test_handles_getMore_error(self):
         client = await self.async_rs_or_single_client(event_listeners=[listener])
 
         collection = client.db["coll"]
-        self.addToCleanup(collection.drop)
+        self.addAsyncCleanup(collection.drop)
         await collection.drop()
 
         fail_command = {
@@ -464,7 +464,7 @@ async def test_no_batch_splits_if_new_namespace_is_not_too_large(self):
                 document={"a": "b"},
             )
         )
-        self.addToCleanup(client.db["coll"].drop)
+        self.addAsyncCleanup(client.db["coll"].drop)
 
         # No batch splitting required.
         result = await client.bulk_write(models=models)
@@ -497,8 +497,8 @@ async def test_batch_splits_if_new_namespace_is_too_large(self):
                 document={"a": "b"},
             )
         )
-        self.addToCleanup(client.db["coll"].drop)
-        self.addToCleanup(client.db[c_repeated].drop)
+        self.addAsyncCleanup(client.db["coll"].drop)
+        self.addAsyncCleanup(client.db[c_repeated].drop)
 
         # Batch splitting required.
         result = await client.bulk_write(models=models)
@@ -561,7 +561,7 @@ async def test_upserted_result(self):
         client = await self.async_rs_or_single_client()
 
         collection = client.db["coll"]
-        self.addToCleanup(collection.drop)
+        self.addAsyncCleanup(collection.drop)
         await collection.drop()
 
         models = []
@@ -602,7 +602,7 @@ async def test_15_unacknowledged_write_across_batches(self):
         client = await self.async_rs_or_single_client(event_listeners=[listener])
 
         collection = client.db["coll"]
-        self.addToCleanup(collection.drop)
+        self.addAsyncCleanup(collection.drop)
         await collection.drop()
         await client.db.command({"create": "db.coll"})
 
@@ -651,10 +651,9 @@ async def test_timeout_in_multi_batch_bulk_write(self):
         _OVERHEAD = 500
 
         internal_client = await self.async_rs_or_single_client(timeoutMS=None)
-        self.addToCleanup(internal_client.close)
 
         collection = internal_client.db["coll"]
-        self.addToCleanup(collection.drop)
+        self.addAsyncCleanup(collection.drop)
         await collection.drop()
 
         fail_command = {