PYTHON-4768 - Fix atlas connection tests and cleanup uses of raw MongoClients in tests (#1867)

test/atlas/test_connection.py

@@ -19,6 +19,7 @@ import os
 import sys
 import unittest
 from collections import defaultdict
+from test import PyMongoTestCase
 
 import pytest
 
@@ -46,38 +47,37 @@ URIS = {
 }
 
 
-def connect(uri):
-    if not uri:
-        raise Exception("Must set env variable to test.")
-    client = pymongo.MongoClient(uri)
-    # No TLS error
-    client.admin.command("ping")
-    # No auth error
-    client.test.test.count_documents({})
+class TestAtlasConnect(PyMongoTestCase):
+    def connect(self, uri):
+        if not uri:
+            raise Exception("Must set env variable to test.")
+        client = self.simple_client(uri)
+        # No TLS error
+        client.admin.command("ping")
+        # No auth error
+        client.test.test.count_documents({})
 
-
-class TestAtlasConnect(unittest.TestCase):
     @unittest.skipUnless(HAS_SNI, "Free tier requires SNI support")
     def test_free_tier(self):
-        connect(URIS["ATLAS_FREE"])
+        self.connect(URIS["ATLAS_FREE"])
 
     def test_replica_set(self):
-        connect(URIS["ATLAS_REPL"])
+        self.connect(URIS["ATLAS_REPL"])
 
     def test_sharded_cluster(self):
-        connect(URIS["ATLAS_SHRD"])
+        self.connect(URIS["ATLAS_SHRD"])
 
     def test_tls_11(self):
-        connect(URIS["ATLAS_TLS11"])
+        self.connect(URIS["ATLAS_TLS11"])
 
     def test_tls_12(self):
-        connect(URIS["ATLAS_TLS12"])
+        self.connect(URIS["ATLAS_TLS12"])
 
     def test_serverless(self):
-        connect(URIS["ATLAS_SERVERLESS"])
+        self.connect(URIS["ATLAS_SERVERLESS"])
 
     def connect_srv(self, uri):
-        connect(uri)
+        self.connect(uri)
         self.assertIn("mongodb+srv://", uri)
 
     @unittest.skipUnless(HAS_SNI, "Free tier requires SNI support")

test/mockupdb/test_auth_recovering_member.py

@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import unittest
+from test import PyMongoTestCase
 
 import pytest
 
@@ -30,7 +31,7 @@ from pymongo.errors import ServerSelectionTimeoutError
 pytestmark = pytest.mark.mockupdb
 
 
-class TestAuthRecoveringMember(unittest.TestCase):
+class TestAuthRecoveringMember(PyMongoTestCase):
     def test_auth_recovering_member(self):
         # Test that we don't attempt auth against a recovering RS member.
         server = MockupDB()
@@ -48,12 +49,10 @@ class TestAuthRecoveringMember(unittest.TestCase):
         server.run()
         self.addCleanup(server.stop)
 
-        client = MongoClient(
+        client = self.simple_client(
             server.uri, replicaSet="rs", serverSelectionTimeoutMS=100, socketTimeoutMS=100
         )
 
-        self.addCleanup(client.close)
-
         # Should see there's no primary or secondary and raise selection timeout
         # error. If it raises AutoReconnect we know it actually tried the
         # server, and that's wrong.

test/mockupdb/test_cluster_time.py

@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import unittest
+from test import PyMongoTestCase
 
 import pytest
 
@@ -34,7 +35,7 @@ from pymongo.errors import OperationFailure
 pytestmark = pytest.mark.mockupdb
 
 
-class TestClusterTime(unittest.TestCase):
+class TestClusterTime(PyMongoTestCase):
     def cluster_time_conversation(self, callback, replies, max_wire_version=6):
         cluster_time = Timestamp(0, 0)
         server = MockupDB()
@@ -52,8 +53,7 @@ class TestClusterTime(unittest.TestCase):
         server.run()
         self.addCleanup(server.stop)
 
-        client = MongoClient(server.uri)
-        self.addCleanup(client.close)
+        client = self.simple_client(server.uri)
 
         with going(callback, client):
             for reply in replies:
@@ -118,8 +118,7 @@ class TestClusterTime(unittest.TestCase):
         server.run()
         self.addCleanup(server.stop)
 
-        client = MongoClient(server.uri, heartbeatFrequencyMS=500)
-        self.addCleanup(client.close)
+        client = self.simple_client(server.uri, heartbeatFrequencyMS=500)
 
         request = server.receives("ismaster")
         # No $clusterTime in first ismaster, only in subsequent ones

test/mockupdb/test_cursor_namespace.py

@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import unittest
+from test import PyMongoTestCase
 
 import pytest
 
@@ -32,7 +33,7 @@ from pymongo import MongoClient
 pytestmark = pytest.mark.mockupdb
 
 
-class TestCursorNamespace(unittest.TestCase):
+class TestCursorNamespace(PyMongoTestCase):
     server: MockupDB
     client: MongoClient
 
@@ -40,7 +41,7 @@ class TestCursorNamespace(unittest.TestCase):
     def setUpClass(cls):
         cls.server = MockupDB(auto_ismaster={"maxWireVersion": 6})
         cls.server.run()
-        cls.client = MongoClient(cls.server.uri)
+        cls.client = cls.unmanaged_simple_client(cls.server.uri)
 
     @classmethod
     def tearDownClass(cls):
@@ -88,7 +89,7 @@ class TestCursorNamespace(unittest.TestCase):
         self._test_cursor_namespace(op, "listIndexes")
 
 
-class TestKillCursorsNamespace(unittest.TestCase):
+class TestKillCursorsNamespace(PyMongoTestCase):
     server: MockupDB
     client: MongoClient
 
@@ -96,7 +97,7 @@ class TestKillCursorsNamespace(unittest.TestCase):
     def setUpClass(cls):
         cls.server = MockupDB(auto_ismaster={"maxWireVersion": 6})
         cls.server.run()
-        cls.client = MongoClient(cls.server.uri)
+        cls.client = cls.unmanaged_simple_client(cls.server.uri)
 
     @classmethod
     def tearDownClass(cls):

test/mockupdb/test_getmore_sharded.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 
 import unittest
 from queue import Queue
+from test import PyMongoTestCase
 
 import pytest
 
@@ -33,7 +34,7 @@ from pymongo import MongoClient
 pytestmark = pytest.mark.mockupdb
 
 
-class TestGetmoreSharded(unittest.TestCase):
+class TestGetmoreSharded(PyMongoTestCase):
     def test_getmore_sharded(self):
         servers = [MockupDB(), MockupDB()]
 
@@ -47,11 +48,10 @@ class TestGetmoreSharded(unittest.TestCase):
             server.run()
             self.addCleanup(server.stop)
 
-        client = MongoClient(
+        client = self.simple_client(
             "mongodb://%s:%d,%s:%d"
             % (servers[0].host, servers[0].port, servers[1].host, servers[1].port)
         )
-        self.addCleanup(client.close)
         collection = client.db.collection
         cursor = collection.find()
         with going(next, cursor):

test/mockupdb/test_initial_ismaster.py

@@ -15,6 +15,7 @@ from __future__ import annotations
 
 import time
 import unittest
+from test import PyMongoTestCase
 
 import pytest
 
@@ -31,15 +32,14 @@ from pymongo import MongoClient
 pytestmark = pytest.mark.mockupdb
 
 
-class TestInitialIsMaster(unittest.TestCase):
+class TestInitialIsMaster(PyMongoTestCase):
     def test_initial_ismaster(self):
         server = MockupDB()
         server.run()
         self.addCleanup(server.stop)
 
         start = time.time()
-        client = MongoClient(server.uri)
-        self.addCleanup(client.close)
+        client = self.simple_client(server.uri)
 
         # A single ismaster is enough for the client to be connected.
         self.assertFalse(client.nodes)

test/mockupdb/test_list_indexes.py

@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import unittest
+from test import PyMongoTestCase
 
 import pytest
 
@@ -28,18 +29,16 @@ except ImportError:
 
 
 from bson import SON
-from pymongo import MongoClient
 
 pytestmark = pytest.mark.mockupdb
 
 
-class TestListIndexes(unittest.TestCase):
+class TestListIndexes(PyMongoTestCase):
     def test_list_indexes_command(self):
         server = MockupDB(auto_ismaster={"maxWireVersion": 6})
         server.run()
         self.addCleanup(server.stop)
-        client = MongoClient(server.uri)
-        self.addCleanup(client.close)
+        client = self.simple_client(server.uri)
         with going(client.test.collection.list_indexes) as cursor:
             request = server.receives(listIndexes="collection", namespace="test")
             request.reply({"cursor": {"firstBatch": [{"name": "index_0"}], "id": 123}})

test/mockupdb/test_max_staleness.py

@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import unittest
+from test import PyMongoTestCase
 
 import pytest
 
@@ -30,7 +31,7 @@ from pymongo import MongoClient
 pytestmark = pytest.mark.mockupdb
 
 
-class TestMaxStalenessMongos(unittest.TestCase):
+class TestMaxStalenessMongos(PyMongoTestCase):
     def test_mongos(self):
         mongos = MockupDB()
         mongos.autoresponds("ismaster", maxWireVersion=6, ismaster=True, msg="isdbgrid")
@@ -40,8 +41,7 @@ class TestMaxStalenessMongos(unittest.TestCase):
         # No maxStalenessSeconds.
         uri = "mongodb://localhost:%d/?readPreference=secondary" % mongos.port
 
-        client = MongoClient(uri)
-        self.addCleanup(client.close)
+        client = self.simple_client(uri)
         with going(client.db.coll.find_one) as future:
             request = mongos.receives()
             self.assertNotIn("maxStalenessSeconds", request.doc["$readPreference"])
@@ -60,8 +60,7 @@ class TestMaxStalenessMongos(unittest.TestCase):
             "&maxStalenessSeconds=1" % mongos.port
         )
 
-        client = MongoClient(uri)
-        self.addCleanup(client.close)
+        client = self.simple_client(uri)
         with going(client.db.coll.find_one) as future:
             request = mongos.receives()
             self.assertEqual(1, request.doc["$readPreference"]["maxStalenessSeconds"])

test/mockupdb/test_mixed_version_sharded.py

@@ -18,6 +18,7 @@ from __future__ import annotations
 import time
 import unittest
 from queue import Queue
+from test import PyMongoTestCase
 
 import pytest
 
@@ -35,7 +36,7 @@ from pymongo import MongoClient
 pytestmark = pytest.mark.mockupdb
 
 
-class TestMixedVersionSharded(unittest.TestCase):
+class TestMixedVersionSharded(PyMongoTestCase):
     def setup_server(self, upgrade):
         self.mongos_old, self.mongos_new = MockupDB(), MockupDB()
 
@@ -62,7 +63,7 @@ class TestMixedVersionSharded(unittest.TestCase):
             self.mongos_new.address_string,
         )
 
-        self.client = MongoClient(self.mongoses_uri)
+        self.client = self.simple_client(self.mongoses_uri)
 
     def tearDown(self):
         if hasattr(self, "client") and self.client:

test/mockupdb/test_op_msg.py

@@ -15,6 +15,7 @@ from __future__ import annotations
 
 import unittest
 from collections import namedtuple
+from test import PyMongoTestCase
 
 import pytest
 
@@ -273,7 +274,7 @@ else:
     operations_312 = []
 
 
-class TestOpMsg(unittest.TestCase):
+class TestOpMsg(PyMongoTestCase):
     server: MockupDB
     client: MongoClient
 
@@ -281,7 +282,7 @@ class TestOpMsg(unittest.TestCase):
     def setUpClass(cls):
         cls.server = MockupDB(auto_ismaster=True, max_wire_version=8)
         cls.server.run()
-        cls.client = MongoClient(cls.server.uri)
+        cls.client = cls.unmanaged_simple_client(cls.server.uri)
 
     @classmethod
     def tearDownClass(cls):

test/mockupdb/test_op_msg_read_preference.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 import copy
 import itertools
 import unittest
+from test import PyMongoTestCase
 from typing import Any
 
 import pytest
@@ -39,7 +40,7 @@ from pymongo.read_preferences import (
 pytestmark = pytest.mark.mockupdb
 
 
-class OpMsgReadPrefBase(unittest.TestCase):
+class OpMsgReadPrefBase(PyMongoTestCase):
     single_mongod = False
     primary: MockupDB
     secondary: MockupDB
@@ -53,8 +54,7 @@ class OpMsgReadPrefBase(unittest.TestCase):
         setattr(cls, test_name, test)
 
     def setup_client(self, read_preference):
-        client = MongoClient(self.primary.uri, read_preference=read_preference)
-        self.addCleanup(client.close)
+        client = self.simple_client(self.primary.uri, read_preference=read_preference)
         return client
 
 
@@ -115,12 +115,13 @@ class TestOpMsgReplicaSet(OpMsgReadPrefBase):
             setattr(cls, test_name, test)
 
     def setup_client(self, read_preference):
-        client = MongoClient(self.primary.uri, replicaSet="rs", read_preference=read_preference)
+        client = self.simple_client(
+            self.primary.uri, replicaSet="rs", read_preference=read_preference
+        )
 
         # Run a command on a secondary to discover the topology. This ensures
         # that secondaryPreferred commands will select the secondary.
         client.admin.command("ismaster", read_preference=ReadPreference.SECONDARY)
-        self.addCleanup(client.close)
         return client
 
 

test/mockupdb/test_query_read_pref_sharded.py

@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import unittest
+from test import PyMongoTestCase
 
 import pytest
 
@@ -40,7 +41,7 @@ from pymongo.read_preferences import (
 pytestmark = pytest.mark.mockupdb
 
 
-class TestQueryAndReadModeSharded(unittest.TestCase):
+class TestQueryAndReadModeSharded(PyMongoTestCase):
     def test_query_and_read_mode_sharded_op_msg(self):
         """Test OP_MSG sends non-primary $readPreference and never $query."""
         server = MockupDB()
@@ -50,8 +51,7 @@ class TestQueryAndReadModeSharded(unittest.TestCase):
         server.run()
         self.addCleanup(server.stop)
 
-        client = MongoClient(server.uri)
-        self.addCleanup(client.close)
+        client = self.simple_client(server.uri)
 
         read_prefs = (
             Primary(),

test/mockupdb/test_reset_and_request_check.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 import itertools
 import time
 import unittest
+from test import PyMongoTestCase
 
 import pytest
 
@@ -37,7 +38,7 @@ from pymongo.server_type import SERVER_TYPE
 pytestmark = pytest.mark.mockupdb
 
 
-class TestResetAndRequestCheck(unittest.TestCase):
+class TestResetAndRequestCheck(PyMongoTestCase):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.ismaster_time = 0.0
@@ -58,7 +59,7 @@ class TestResetAndRequestCheck(unittest.TestCase):
         kwargs = {"socketTimeoutMS": 100}
         # Disable retryable reads when pymongo supports it.
         kwargs["retryReads"] = False
-        self.client = MongoClient(self.server.uri, **kwargs)  # type: ignore
+        self.client = self.simple_client(self.server.uri, **kwargs)  # type: ignore
         wait_until(lambda: self.client.nodes, "connect to standalone")
 
     def tearDown(self):

test/test_index_management.py

@@ -25,11 +25,10 @@ import pytest
 
 sys.path[0:0] = [""]
 
-from test import IntegrationTest, unittest
+from test import IntegrationTest, PyMongoTestCase, unittest
 from test.unified_format import generate_test_classes
 from test.utils import AllowListEventListener, EventListener
 
-from pymongo import MongoClient
 from pymongo.errors import OperationFailure
 from pymongo.operations import SearchIndexModel
 from pymongo.read_concern import ReadConcern
@@ -47,8 +46,7 @@ class TestCreateSearchIndex(IntegrationTest):
         if not os.environ.get("TEST_INDEX_MANAGEMENT"):
             raise unittest.SkipTest("Skipping index management tests")
         listener = AllowListEventListener("createSearchIndexes")
-        client = MongoClient(event_listeners=[listener])
-        self.addCleanup(client.close)
+        client = self.simple_client(event_listeners=[listener])
         coll = client.test.test
         coll.drop()
         definition = dict(mappings=dict(dynamic=True))
@@ -79,7 +77,7 @@ class TestCreateSearchIndex(IntegrationTest):
         )
 
 
-class SearchIndexIntegrationBase(unittest.TestCase):
+class SearchIndexIntegrationBase(PyMongoTestCase):
     db_name = "test_search_index_base"
 
     @classmethod
@@ -91,7 +89,7 @@ class SearchIndexIntegrationBase(unittest.TestCase):
         username = os.environ["DB_USER"]
         password = os.environ["DB_PASSWORD"]
         cls.listener = listener = EventListener()
-        cls.client = MongoClient(
+        cls.client = cls.unmanaged_simple_client(
             url, username=username, password=password, event_listeners=[listener]
         )
         cls.client.drop_database(_NAME)

test/test_retryable_reads.py

@@ -43,7 +43,6 @@ from pymongo.monitoring import (
     ConnectionCheckOutFailedReason,
     PoolClearedEvent,
 )
-from pymongo.synchronous.mongo_client import MongoClient
 
 # Location of JSON test specifications.
 _TEST_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "retryable_reads", "legacy")
@@ -51,19 +50,19 @@ _TEST_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "retryabl
 
 class TestClientOptions(PyMongoTestCase):
     def test_default(self):
-        client = MongoClient(connect=False)
+        client = self.simple_client(connect=False)
         self.assertEqual(client.options.retry_reads, True)
 
     def test_kwargs(self):
-        client = MongoClient(retryReads=True, connect=False)
+        client = self.simple_client(retryReads=True, connect=False)
         self.assertEqual(client.options.retry_reads, True)
-        client = MongoClient(retryReads=False, connect=False)
+        client = self.simple_client(retryReads=False, connect=False)
         self.assertEqual(client.options.retry_reads, False)
 
     def test_uri(self):
-        client = MongoClient("mongodb://h/?retryReads=true", connect=False)
+        client = self.simple_client("mongodb://h/?retryReads=true", connect=False)
         self.assertEqual(client.options.retry_reads, True)
-        client = MongoClient("mongodb://h/?retryReads=false", connect=False)
+        client = self.simple_client("mongodb://h/?retryReads=false", connect=False)
         self.assertEqual(client.options.retry_reads, False)