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

Author: NoahStapp

test/atlas/test_connection.py

@@ -19,6 +19,7 @@
 import sys
 import unittest
 from collections import defaultdict
+from test import PyMongoTestCase
 
 import pytest
 
@@ -46,38 +47,37 @@
 }
 
 
-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 @@
 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 @@ def test_auth_recovering_member(self):
         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 @@
 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 @@ def cluster_time_conversation(self, callback, replies, max_wire_version=6):
         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 @@ def test_monitor(self):
         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,15 +33,15 @@
 pytestmark = pytest.mark.mockupdb
 
 
-class TestCursorNamespace(unittest.TestCase):
+class TestCursorNamespace(PyMongoTestCase):
     server: MockupDB
     client: MongoClient
 
     @classmethod
     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,15 +89,15 @@ def op():
         self._test_cursor_namespace(op, "listIndexes")
 
 
-class TestKillCursorsNamespace(unittest.TestCase):
+class TestKillCursorsNamespace(PyMongoTestCase):
     server: MockupDB
     client: MongoClient
 
     @classmethod
     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 @@
 
 import unittest
 from queue import Queue
+from test import PyMongoTestCase
 
 import pytest
 
@@ -33,7 +34,7 @@
 pytestmark = pytest.mark.mockupdb
 
 
-class TestGetmoreSharded(unittest.TestCase):
+class TestGetmoreSharded(PyMongoTestCase):
     def test_getmore_sharded(self):
         servers = [MockupDB(), MockupDB()]
 
@@ -47,11 +48,10 @@ def test_getmore_sharded(self):
             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 @@
 
 import time
 import unittest
+from test import PyMongoTestCase
 
 import pytest
 
@@ -31,15 +32,14 @@
 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 @@
 
 
 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 @@
 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 @@ def test_mongos(self):
         # 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 @@ def test_mongos(self):
             "&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 @@
 import time
 import unittest
 from queue import Queue
+from test import PyMongoTestCase
 
 import pytest
 
@@ -35,7 +36,7 @@
 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 @@ def setup_server(self, upgrade):
             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 @@
 
 import unittest
 from collections import namedtuple
+from test import PyMongoTestCase
 
 import pytest
 
@@ -273,15 +274,15 @@
     operations_312 = []
 
 
-class TestOpMsg(unittest.TestCase):
+class TestOpMsg(PyMongoTestCase):
     server: MockupDB
     client: MongoClient
 
     @classmethod
     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 @@
 import copy
 import itertools
 import unittest
+from test import PyMongoTestCase
 from typing import Any
 
 import pytest
@@ -39,7 +40,7 @@
 pytestmark = pytest.mark.mockupdb
 
 
-class OpMsgReadPrefBase(unittest.TestCase):
+class OpMsgReadPrefBase(PyMongoTestCase):
     single_mongod = False
     primary: MockupDB
     secondary: MockupDB
@@ -53,8 +54,7 @@ def add_test(cls, mode, test_name, test):
         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 @@ def add_test(cls, mode, test_name, test):
             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