Convert test_client to pytest

NoahStapp · NoahStapp · commit 266b0a3b150c · 2025-01-23T14:37:37.000-05:00
diff --git a/pymongo/asynchronous/mongo_client.py b/pymongo/asynchronous/mongo_client.py
@@ -1418,6 +1418,7 @@ def __next__(self) -> NoReturn:
         raise TypeError("'AsyncMongoClient' object is not iterable")
 
     next = __next__
+    anext = next
 
     async def _server_property(self, attr_name: str) -> Any:
         """An attribute of the current server's description.
diff --git a/pyproject.toml b/pyproject.toml
@@ -94,7 +94,7 @@ zstd = ["requirements/zstd.txt"]
 
 [tool.pytest.ini_options]
 minversion = "7"
-addopts = ["-ra", "--strict-config", "--strict-markers", "--junitxml=xunit-results/TEST-results.xml", "-m default or default_async"]
+addopts = ["-ra", "--strict-config", "--strict-markers", "--junitxml=xunit-results/TEST-results.xml", "-m default or default_async or asyncio"]
 testpaths = ["test"]
 log_cli_level = "INFO"
 faulthandler_timeout = 1500
diff --git a/test/asynchronous/__init__.py b/test/asynchronous/__init__.py
@@ -30,6 +30,10 @@
 import unittest
 import warnings
 from asyncio import iscoroutinefunction
+
+import pytest_asyncio
+from pymongo.lock import _create_lock, _async_create_lock
+
 from test.helpers import (
     COMPRESSORS,
     IS_SRV,
@@ -116,7 +120,7 @@ def __init__(self):
         self.default_client_options: Dict = {}
         self.sessions_enabled = False
         self.client = None  # type: ignore
-        self.conn_lock = threading.Lock()
+        self.conn_lock = _async_create_lock()
         self.is_data_lake = False
         self.load_balancer = TEST_LOADBALANCER
         self.serverless = TEST_SERVERLESS
@@ -337,7 +341,7 @@ async def _init_client(self):
                         await mongos_client.close()
 
     async def init(self):
-        with self.conn_lock:
+        async with self.conn_lock:
             if not self.client and not self.connection_attempts:
                 await self._init_client()
 
@@ -520,6 +524,12 @@ def require_data_lake(self, func):
             func=func,
         )
 
+    @property
+    def is_not_mmap(self):
+        if self.is_mongos:
+            return True
+        return self.storage_engine != "mmapv1"
+
     def require_no_mmap(self, func):
         """Run a test only if the server is not using the MMAPv1 storage
         engine. Only works for standalone and replica sets; tests are
@@ -573,6 +583,10 @@ def require_replica_set(self, func):
         """Run a test only if the client is connected to a replica set."""
         return self._require(lambda: self.is_rs, "Not connected to a replica set", func=func)
 
+    @property
+    async def secondaries_count(self):
+        return 0 if not self.client else len(await self.client.secondaries)
+
     def require_secondaries_count(self, count):
         """Run a test only if the client is connected to a replica set that has
         `count` secondaries.
@@ -588,7 +602,7 @@ async def check():
 
     @property
     async def supports_secondary_read_pref(self):
-        if self.has_secondaries:
+        if await self.has_secondaries:
             return True
         if self.is_mongos:
             shard = await self.client.config.shards.find_one()["host"]  # type:ignore[index]
@@ -692,7 +706,7 @@ async def is_topology_type(self, topologies):
         if "sharded" in topologies and self.is_mongos:
             return True
         if "sharded-replicaset" in topologies and self.is_mongos:
-            shards = await async_client_context.client.config.shards.find().to_list()
+            shards = await self.client.config.shards.find().to_list()
             for shard in shards:
                 # For a 3-member RS-backed sharded cluster, shard['host']
                 # will be 'replicaName/ip1:port1,ip2:port2,ip3:port3'
@@ -1191,8 +1205,39 @@ async def asyncTearDown(self) -> None:
         await super().asyncTearDown()
 
 
+async def _get_environment():
+    client = AsyncClientContext()
+    await client.init()
+    requirements = {}
+    requirements["SUPPORT_TRANSACTIONS"] = client.supports_transactions()
+    requirements["IS_DATA_LAKE"] = client.is_data_lake
+    requirements["IS_SYNC"] = _IS_SYNC
+    requirements["IS_SYNC"] = _IS_SYNC
+    requirements["REQUIRE_API_VERSION"] = MONGODB_API_VERSION
+    requirements["SUPPORTS_FAILCOMMAND_FAIL_POINT"] = client.supports_failCommand_fail_point
+    requirements["IS_NOT_MMAP"] = client.is_not_mmap
+    requirements["SERVER_VERSION"] = client.version
+    requirements["AUTH_ENABLED"] = client.auth_enabled
+    requirements["FIPS_ENABLED"] = client.fips_enabled
+    requirements["IS_RS"] = client.is_rs
+    requirements["MONGOSES"] = len(client.mongoses)
+    requirements["SECONDARIES_COUNT"] = await client.secondaries_count
+    requirements["SECONDARY_READ_PREF"] = await client.supports_secondary_read_pref
+    requirements["HAS_IPV6"] = client.has_ipv6
+    requirements["IS_SERVERLESS"] = client.serverless
+    requirements["IS_LOAD_BALANCER"] = client.load_balancer
+    requirements["TEST_COMMANDS_ENABLED"] = client.test_commands_enabled
+    requirements["IS_TLS"] = client.tls
+    requirements["IS_TLS_CERT"] = client.tlsCertificateKeyFile
+    requirements["SERVER_IS_RESOLVEABLE"] = client.server_is_resolvable
+    requirements["SESSIONS_ENABLED"] = client.sessions_enabled
+    requirements["SUPPORTS_RETRYABLE_WRITES"] = client.supports_retryable_writes()
+    await client.client.close()
+
+    return requirements
+
 async def async_setup():
-    await async_client_context.init()
+    await _get_environment()
     warnings.resetwarnings()
     warnings.simplefilter("always")
     global_knobs.enable()
@@ -1207,16 +1252,16 @@ async def async_teardown():
         garbage.append(f"  gc.get_referrers: {gc.get_referrers(g)!r}")
     if garbage:
         raise AssertionError("\n".join(garbage))
-    c = async_client_context.client
-    if c:
-        if not async_client_context.is_data_lake:
-            await c.drop_database("pymongo-pooling-tests")
-            await c.drop_database("pymongo_test")
-            await c.drop_database("pymongo_test1")
-            await c.drop_database("pymongo_test2")
-            await c.drop_database("pymongo_test_mike")
-            await c.drop_database("pymongo_test_bernie")
-        await c.close()
+    # c = async_client_context.client
+    # if c:
+    #     if not async_client_context.is_data_lake:
+    #         await c.drop_database("pymongo-pooling-tests")
+    #         await c.drop_database("pymongo_test")
+    #         await c.drop_database("pymongo_test1")
+    #         await c.drop_database("pymongo_test2")
+    #         await c.drop_database("pymongo_test_mike")
+    #         await c.drop_database("pymongo_test_bernie")
+    #     await c.close()
     print_running_clients()
 
 
diff --git a/test/asynchronous/conftest.py b/test/asynchronous/conftest.py
@@ -2,8 +2,14 @@
 
 import asyncio
 import sys
-from test import pytest_conf
-from test.asynchronous import async_setup, async_teardown
+
+from typing_extensions import Any
+
+from pymongo import AsyncMongoClient
+from pymongo.uri_parser import parse_uri
+
+from test import pytest_conf, db_user, db_pwd
+from test.asynchronous import async_setup, async_teardown, _connection_string, AsyncClientContext
 
 import pytest
 import pytest_asyncio
@@ -21,12 +27,103 @@ def event_loop_policy():
 
     return asyncio.get_event_loop_policy()
 
+@pytest_asyncio.fixture(loop_scope="session")
+async def async_client_context_fixture():
+    client = AsyncClientContext()
+    await client.init()
+    yield client
+    await client.client.close()
 
-@pytest_asyncio.fixture(scope="package", autouse=True)
+@pytest_asyncio.fixture(loop_scope="session", autouse=True)
 async def test_setup_and_teardown():
     await async_setup()
     yield
     await async_teardown()
 
+async def _async_mongo_client(
+        async_client_context, host, port, authenticate=True, directConnection=None, **kwargs
+    ):
+        """Create a new client over SSL/TLS if necessary."""
+        host = host or await async_client_context.host
+        port = port or await async_client_context.port
+        client_options: dict = async_client_context.default_client_options.copy()
+        if async_client_context.replica_set_name and not directConnection:
+            client_options["replicaSet"] = async_client_context.replica_set_name
+        if directConnection is not None:
+            client_options["directConnection"] = directConnection
+        client_options.update(kwargs)
+
+        uri = _connection_string(host)
+        auth_mech = kwargs.get("authMechanism", "")
+        if async_client_context.auth_enabled and authenticate and auth_mech != "MONGODB-OIDC":
+            # Only add the default username or password if one is not provided.
+            res = parse_uri(uri)
+            if (
+                not res["username"]
+                and not res["password"]
+                and "username" not in client_options
+                and "password" not in client_options
+            ):
+                client_options["username"] = db_user
+                client_options["password"] = db_pwd
+        client = AsyncMongoClient(uri, port, **client_options)
+        if client._options.connect:
+            await client.aconnect()
+        return client
+
+
+async def async_single_client_noauth(
+    async_client_context, h: Any = None, p: Any = None, **kwargs: Any
+) -> AsyncMongoClient[dict]:
+    """Make a direct connection. Don't authenticate."""
+    return await _async_mongo_client(async_client_context, h, p, authenticate=False, directConnection=True, **kwargs)
+#
+async def async_single_client(
+    async_client_context, h: Any = None, p: Any = None, **kwargs: Any
+) -> AsyncMongoClient[dict]:
+    """Make a direct connection, and authenticate if necessary."""
+    return await _async_mongo_client(async_client_context, h, p, directConnection=True, **kwargs)
+
+# @pytest_asyncio.fixture(loop_scope="function")
+# async def async_rs_client_noauth(
+#     async_client_context, h: Any = None, p: Any = None, **kwargs: Any
+# ) -> AsyncMongoClient[dict]:
+#     """Connect to the replica set. Don't authenticate."""
+#     return await _async_mongo_client(async_client_context, h, p, authenticate=False, **kwargs)
+#
+# @pytest_asyncio.fixture(loop_scope="function")
+# async def async_rs_client(
+#     async_client_context, h: Any = None, p: Any = None, **kwargs: Any
+# ) -> AsyncMongoClient[dict]:
+#     """Connect to the replica set and authenticate if necessary."""
+#     return await _async_mongo_client(async_client_context, h, p, **kwargs)
+#
+# @pytest_asyncio.fixture(loop_scope="function")
+# async def async_rs_or_single_client_noauth(
+#     async_client_context, h: Any = None, p: Any = None, **kwargs: Any
+# ) -> AsyncMongoClient[dict]:
+#     """Connect to the replica set if there is one, otherwise the standalone.
+#
+#     Like rs_or_single_client, but does not authenticate.
+#     """
+#     return await _async_mongo_client(async_client_context, h, p, authenticate=False, **kwargs)
+
+async def async_rs_or_single_client(
+    async_client_context, h: Any = None, p: Any = None, **kwargs: Any
+) -> AsyncMongoClient[Any]:
+    """Connect to the replica set if there is one, otherwise the standalone.
+
+    Authenticates if necessary.
+    """
+    return await _async_mongo_client(async_client_context, h, p, **kwargs)
+
+def simple_client(h: Any = None, p: Any = None, **kwargs: Any) -> AsyncMongoClient:
+    if not h and not p:
+        client = AsyncMongoClient(**kwargs)
+    else:
+        client = AsyncMongoClient(h, p, **kwargs)
+    return client
+
+
 
 pytest_collection_modifyitems = pytest_conf.pytest_collection_modifyitems
diff --git a/test/asynchronous/test_client_pytest.py b/test/asynchronous/test_client_pytest.py
diff --git a/test/pytest_conf.py b/test/pytest_conf.py
