Fix tests

Author: NoahStapp

pymongo/asynchronous/auth_oidc.py

@@ -15,7 +15,7 @@
 """MONGODB-OIDC Authentication helpers."""
 from __future__ import annotations
 
-import threading
+import asyncio
 import time
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, Optional, Union
@@ -36,6 +36,7 @@
 )
 from pymongo.errors import ConfigurationError, OperationFailure
 from pymongo.helpers_shared import _AUTHENTICATION_FAILURE_CODE
+from pymongo.lock import Lock, _async_create_lock
 
 if TYPE_CHECKING:
     from pymongo.asynchronous.pool import AsyncConnection
@@ -81,7 +82,7 @@ class _OIDCAuthenticator:
     access_token: Optional[str] = field(default=None)
     idp_info: Optional[OIDCIdPInfo] = field(default=None)
     token_gen_id: int = field(default=0)
-    lock: threading.Lock = field(default_factory=threading.Lock)
+    lock: Lock = field(default_factory=_async_create_lock)
     last_call_time: float = field(default=0)
 
     async def reauthenticate(self, conn: AsyncConnection) -> Optional[Mapping[str, Any]]:
@@ -164,7 +165,7 @@ async def _authenticate_human(self, conn: AsyncConnection) -> Optional[Mapping[s
         # Attempt to authenticate with a JwtStepRequest.
         return await self._sasl_continue_jwt(conn, start_resp)
 
-    def _get_access_token(self) -> Optional[str]:
+    async def _get_access_token(self) -> Optional[str]:
         properties = self.properties
         cb: Union[None, OIDCCallback]
         resp: OIDCCallbackResult
@@ -186,7 +187,7 @@ def _get_access_token(self) -> Optional[str]:
             return None
 
         if not prev_token and cb is not None:
-            with self.lock:
+            async with self.lock:
                 # See if the token was changed while we were waiting for the
                 # lock.
                 new_token = self.access_token
@@ -196,7 +197,7 @@ def _get_access_token(self) -> Optional[str]:
                 # Ensure that we are waiting a min time between callback invocations.
                 delta = time.time() - self.last_call_time
                 if delta < TIME_BETWEEN_CALLS_SECONDS:
-                    time.sleep(TIME_BETWEEN_CALLS_SECONDS - delta)
+                    await asyncio.sleep(TIME_BETWEEN_CALLS_SECONDS - delta)
                 self.last_call_time = time.time()
 
                 if is_human:
@@ -211,7 +212,10 @@ def _get_access_token(self) -> Optional[str]:
                     idp_info=self.idp_info,
                     username=self.properties.username,
                 )
-                resp = cb.fetch(context)
+                if not _IS_SYNC:
+                    resp = await asyncio.get_running_loop().run_in_executor(None, cb.fetch, context)
+                else:
+                    resp = cb.fetch(context)
                 if not isinstance(resp, OIDCCallbackResult):
                     raise ValueError(
                         f"Callback result must be of type OIDCCallbackResult, not {type(resp)}"
@@ -253,13 +257,13 @@ async def _sasl_continue_jwt(
         start_payload: dict = bson.decode(start_resp["payload"])
         if "issuer" in start_payload:
             self.idp_info = OIDCIdPInfo(**start_payload)
-        access_token = self._get_access_token()
+        access_token = await self._get_access_token()
         conn.oidc_token_gen_id = self.token_gen_id
         cmd = self._get_continue_command({"jwt": access_token}, start_resp)
         return await self._run_command(conn, cmd)
 
     async def _sasl_start_jwt(self, conn: AsyncConnection) -> Mapping[str, Any]:
-        access_token = self._get_access_token()
+        access_token = await self._get_access_token()
         conn.oidc_token_gen_id = self.token_gen_id
         cmd = self._get_start_command({"jwt": access_token})
         return await self._run_command(conn, cmd)

pymongo/asynchronous/cursor.py

@@ -1130,7 +1130,6 @@ async def _send_message(self, operation: Union[_Query, _GetMore]) -> None:
         except BaseException:
             await self.close()
             raise
-
         self._address = response.address
         if isinstance(response, PinnedResponse):
             if not self._sock_mgr:

pymongo/asynchronous/helpers.py

@@ -64,7 +64,7 @@ async def inner(*args: Any, **kwargs: Any) -> Any:
                     await conn.authenticate(reauthenticate=True)
                 else:
                     raise
-                return func(*args, **kwargs)
+                return await func(*args, **kwargs)
             raise
 
     return cast(F, inner)

pymongo/synchronous/auth_oidc.py

@@ -15,7 +15,7 @@
 """MONGODB-OIDC Authentication helpers."""
 from __future__ import annotations
 
-import threading
+import asyncio
 import time
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, Optional, Union
@@ -36,6 +36,7 @@
 )
 from pymongo.errors import ConfigurationError, OperationFailure
 from pymongo.helpers_shared import _AUTHENTICATION_FAILURE_CODE
+from pymongo.lock import Lock, _create_lock
 
 if TYPE_CHECKING:
     from pymongo.auth_shared import MongoCredential
@@ -81,7 +82,7 @@ class _OIDCAuthenticator:
     access_token: Optional[str] = field(default=None)
     idp_info: Optional[OIDCIdPInfo] = field(default=None)
     token_gen_id: int = field(default=0)
-    lock: threading.Lock = field(default_factory=threading.Lock)
+    lock: Lock = field(default_factory=_create_lock)
     last_call_time: float = field(default=0)
 
     def reauthenticate(self, conn: Connection) -> Optional[Mapping[str, Any]]:
@@ -211,7 +212,10 @@ def _get_access_token(self) -> Optional[str]:
                     idp_info=self.idp_info,
                     username=self.properties.username,
                 )
-                resp = cb.fetch(context)
+                if not _IS_SYNC:
+                    resp = asyncio.get_running_loop().run_in_executor(None, cb.fetch, context)
+                else:
+                    resp = cb.fetch(context)
                 if not isinstance(resp, OIDCCallbackResult):
                     raise ValueError(
                         f"Callback result must be of type OIDCCallbackResult, not {type(resp)}"

pymongo/synchronous/cursor.py

@@ -1128,7 +1128,6 @@ def _send_message(self, operation: Union[_Query, _GetMore]) -> None:
         except BaseException:
             self.close()
             raise
-
         self._address = response.address
         if isinstance(response, PinnedResponse):
             if not self._sock_mgr:

test/asynchronous/test_auth_oidc.py

@@ -24,13 +24,14 @@
 from contextlib import asynccontextmanager
 from pathlib import Path
 from test.asynchronous import AsyncPyMongoTestCase
+from test.asynchronous.helpers import ConcurrentRunner
 from typing import Dict
 
 import pytest
 
 sys.path[0:0] = [""]
 
-from test.unified_format import generate_test_classes
+from test.asynchronous.unified_format import generate_test_classes
 from test.utils_shared import EventListener, OvertCommandListener
 
 from bson import SON
@@ -258,10 +259,11 @@ async def test_1_7_allowed_hosts_in_connection_string_ignored(self):
         uri = "mongodb+srv://example.com?authMechanism=MONGODB-OIDC&authMechanismProperties=ALLOWED_HOSTS:%5B%22example.com%22%5D"
         with self.assertRaises(ConfigurationError), warnings.catch_warnings():
             warnings.simplefilter("ignore")
-            _ = AsyncMongoClient(
+            c = AsyncMongoClient(
                 uri,
                 authmechanismproperties=dict(OIDC_HUMAN_CALLBACK=self.create_request_cb()),
             )
+            await c.aconnect()
 
     async def test_1_8_machine_idp_human_callback(self):
         if not os.environ.get("OIDC_IS_LOCAL"):
@@ -752,7 +754,7 @@ async def test_reauthenticate_succeeds_command(self):
 class TestAuthOIDCMachine(OIDCTestBase):
     uri: str
 
-    def asyncSetUp(self):
+    async def asyncSetUp(self):
         self.request_called = 0
 
     def create_request_cb(self, username=None, sleep=0):
@@ -791,24 +793,24 @@ async def test_1_1_callback_is_called_during_reauthentication(self):
         # Assert that the callback was called 1 time.
         self.assertEqual(self.request_called, 1)
 
-    # TODO REPLACE THREADS WITH TASKS
     async def test_1_2_callback_is_called_once_for_multiple_connections(self):
         # Create a ``AsyncMongoClient`` configured with a custom OIDC callback that
         # implements the provider logic.
         client = await self.create_client()
+        await client.aconnect()
 
         # Start 10 threads and run 100 find operations in each thread that all succeed.
         async def target():
             for _ in range(100):
                 await client.test.test.find_one()
 
-        threads = []
-        for _ in range(10):
-            thread = threading.Thread(target=target)
-            thread.start()
-            threads.append(thread)
-        for thread in threads:
-            thread.join()
+        tasks = []
+        for i in range(10):
+            tasks.append(ConcurrentRunner(target=target))
+        for t in tasks:
+            await t.start()
+        for t in tasks:
+            await t.join()
         # Assert that the callback was called 1 time.
         self.assertEqual(self.request_called, 1)
 
@@ -1043,6 +1045,7 @@ async def test_4_4_speculative_authentication_should_be_ignored_on_reauthenticat
         # Create an OIDC configured client that can listen for `SaslStart` commands.
         listener = EventListener()
         client = await self.create_client(event_listeners=[listener])
+        await client.aconnect()
 
         # Preload the *Client Cache* with a valid access token to enforce Speculative Authentication.
         client2 = await self.create_client()
@@ -1107,6 +1110,7 @@ async def test_speculative_auth_success(self):
         client1 = await self.create_client()
         await client1.test.test.find_one()
         client2 = await self.create_client()
+        await client2.aconnect()
 
         # Prime the cache of the second client.
         client2.options.pool_options._credentials.cache.data = (

test/test_auth_oidc.py

@@ -24,6 +24,7 @@
 from contextlib import contextmanager
 from pathlib import Path
 from test import PyMongoTestCase
+from test.helpers import ConcurrentRunner
 from typing import Dict
 
 import pytest
@@ -256,10 +257,11 @@ def test_1_7_allowed_hosts_in_connection_string_ignored(self):
         uri = "mongodb+srv://example.com?authMechanism=MONGODB-OIDC&authMechanismProperties=ALLOWED_HOSTS:%5B%22example.com%22%5D"
         with self.assertRaises(ConfigurationError), warnings.catch_warnings():
             warnings.simplefilter("ignore")
-            _ = MongoClient(
+            c = MongoClient(
                 uri,
                 authmechanismproperties=dict(OIDC_HUMAN_CALLBACK=self.create_request_cb()),
             )
+            c._connect()
 
     def test_1_8_machine_idp_human_callback(self):
         if not os.environ.get("OIDC_IS_LOCAL"):
@@ -789,24 +791,24 @@ def test_1_1_callback_is_called_during_reauthentication(self):
         # Assert that the callback was called 1 time.
         self.assertEqual(self.request_called, 1)
 
-    # TODO REPLACE THREADS WITH TASKS
     def test_1_2_callback_is_called_once_for_multiple_connections(self):
         # Create a ``MongoClient`` configured with a custom OIDC callback that
         # implements the provider logic.
         client = self.create_client()
+        client._connect()
 
         # Start 10 threads and run 100 find operations in each thread that all succeed.
         def target():
             for _ in range(100):
                 client.test.test.find_one()
 
-        threads = []
-        for _ in range(10):
-            thread = threading.Thread(target=target)
-            thread.start()
-            threads.append(thread)
-        for thread in threads:
-            thread.join()
+        tasks = []
+        for i in range(10):
+            tasks.append(ConcurrentRunner(target=target))
+        for t in tasks:
+            t.start()
+        for t in tasks:
+            t.join()
         # Assert that the callback was called 1 time.
         self.assertEqual(self.request_called, 1)
 
@@ -1041,6 +1043,7 @@ def test_4_4_speculative_authentication_should_be_ignored_on_reauthentication(se
         # Create an OIDC configured client that can listen for `SaslStart` commands.
         listener = EventListener()
         client = self.create_client(event_listeners=[listener])
+        client._connect()
 
         # Preload the *Client Cache* with a valid access token to enforce Speculative Authentication.
         client2 = self.create_client()
@@ -1105,6 +1108,7 @@ def test_speculative_auth_success(self):
         client1 = self.create_client()
         client1.test.test.find_one()
         client2 = self.create_client()
+        client2._connect()
 
         # Prime the cache of the second client.
         client2.options.pool_options._credentials.cache.data = (