PYTHON-4768 - Fix atlas connection tests and cleanup uses of raw Mong…

…oClients in tests (#1867)

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