diff --git a/test/atlas/test_connection.py b/test/atlas/test_connection.py index 762ac30115..4dcbba6d11 100644 --- a/test/atlas/test_connection.py +++ b/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") diff --git a/test/mockupdb/test_auth_recovering_member.py b/test/mockupdb/test_auth_recovering_member.py index 95a83a11a7..6eadafaf38 100644 --- a/test/mockupdb/test_auth_recovering_member.py +++ b/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. diff --git a/test/mockupdb/test_cluster_time.py b/test/mockupdb/test_cluster_time.py index 9794843175..761415951b 100644 --- a/test/mockupdb/test_cluster_time.py +++ b/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 diff --git a/test/mockupdb/test_cursor_namespace.py b/test/mockupdb/test_cursor_namespace.py index 57a98373fc..455a3a923b 100644 --- a/test/mockupdb/test_cursor_namespace.py +++ b/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 @@ 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 @@ def op(): 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): diff --git a/test/mockupdb/test_getmore_sharded.py b/test/mockupdb/test_getmore_sharded.py index cf26b10a15..8ba291e4a7 100644 --- a/test/mockupdb/test_getmore_sharded.py +++ b/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): diff --git a/test/mockupdb/test_initial_ismaster.py b/test/mockupdb/test_initial_ismaster.py index 8046f08db5..3eae98716b 100644 --- a/test/mockupdb/test_initial_ismaster.py +++ b/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) diff --git a/test/mockupdb/test_list_indexes.py b/test/mockupdb/test_list_indexes.py index ce6db9e104..ff3363664b 100644 --- a/test/mockupdb/test_list_indexes.py +++ b/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}}) diff --git a/test/mockupdb/test_max_staleness.py b/test/mockupdb/test_max_staleness.py index 40cf7ef00d..7275aaf44b 100644 --- a/test/mockupdb/test_max_staleness.py +++ b/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"]) diff --git a/test/mockupdb/test_mixed_version_sharded.py b/test/mockupdb/test_mixed_version_sharded.py index 72b42deacc..99d428b5d9 100644 --- a/test/mockupdb/test_mixed_version_sharded.py +++ b/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: diff --git a/test/mockupdb/test_op_msg.py b/test/mockupdb/test_op_msg.py index 776d1644dd..4b85c5a48a 100644 --- a/test/mockupdb/test_op_msg.py +++ b/test/mockupdb/test_op_msg.py @@ -15,6 +15,7 @@ import unittest from collections import namedtuple +from test import PyMongoTestCase import pytest @@ -273,7 +274,7 @@ 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): diff --git a/test/mockupdb/test_op_msg_read_preference.py b/test/mockupdb/test_op_msg_read_preference.py index c7c7037f23..86293d0c09 100644 --- a/test/mockupdb/test_op_msg_read_preference.py +++ b/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 diff --git a/test/mockupdb/test_query_read_pref_sharded.py b/test/mockupdb/test_query_read_pref_sharded.py index 6276ee7789..676e71b711 100644 --- a/test/mockupdb/test_query_read_pref_sharded.py +++ b/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 @@ 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 @@ def test_query_and_read_mode_sharded_op_msg(self): server.run() self.addCleanup(server.stop) - client = MongoClient(server.uri) - self.addCleanup(client.close) + client = self.simple_client(server.uri) read_prefs = ( Primary(), diff --git a/test/mockupdb/test_reset_and_request_check.py b/test/mockupdb/test_reset_and_request_check.py index 2919025f09..dd6ad46b1e 100644 --- a/test/mockupdb/test_reset_and_request_check.py +++ b/test/mockupdb/test_reset_and_request_check.py @@ -16,6 +16,7 @@ import itertools import time import unittest +from test import PyMongoTestCase import pytest @@ -37,7 +38,7 @@ 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 @@ def responder(request): 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): diff --git a/test/test_index_management.py b/test/test_index_management.py index 426d0c7349..ec1e363737 100644 --- a/test/test_index_management.py +++ b/test/test_index_management.py @@ -25,11 +25,10 @@ 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 @@ def test_inputs(self): 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 @@ def test_inputs(self): ) -class SearchIndexIntegrationBase(unittest.TestCase): +class SearchIndexIntegrationBase(PyMongoTestCase): db_name = "test_search_index_base" @classmethod @@ -91,7 +89,7 @@ def setUpClass(cls) -> None: 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) diff --git a/test/test_retryable_reads.py b/test/test_retryable_reads.py index a1c72bb7b6..b4fafe4652 100644 --- a/test/test_retryable_reads.py +++ b/test/test_retryable_reads.py @@ -43,7 +43,6 @@ 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 @@ 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)