PYTHON-4700 - Convert CSFLE tests to async (mongodb#1907)

Author: NoahStapp

.evergreen/run-tests.sh

@@ -257,9 +257,9 @@ if [ -z "$GREEN_FRAMEWORK" ]; then
     # Use --capture=tee-sys so pytest prints test output inline:
     # https://docs.pytest.org/en/stable/how-to/capture-stdout-stderr.html
     if [ -z "$TEST_SUITES" ]; then
-      python -m pytest -v --capture=tee-sys --durations=5 --maxfail=10 $TEST_ARGS
+      python -m pytest -v --capture=tee-sys --durations=5 $TEST_ARGS
     else
-      python -m pytest -v --capture=tee-sys --durations=5 --maxfail=10 -m $TEST_SUITES $TEST_ARGS
+      python -m pytest -v --capture=tee-sys --durations=5 -m $TEST_SUITES $TEST_ARGS
     fi
 else
     python green_framework_test.py $GREEN_FRAMEWORK -v $TEST_ARGS

pymongo/asynchronous/encryption.py

@@ -180,10 +180,20 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                 while kms_context.bytes_needed > 0:
                     # CSOT: update timeout.
                     conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
-                    data = conn.recv(kms_context.bytes_needed)
+                    if _IS_SYNC:
+                        data = conn.recv(kms_context.bytes_needed)
+                    else:
+                        from pymongo.network_layer import (  # type: ignore[attr-defined]
+                            async_receive_data_socket,
+                        )
+
+                        data = await async_receive_data_socket(conn, kms_context.bytes_needed)
                     if not data:
                         raise OSError("KMS connection closed")
                     kms_context.feed(data)
+            # Async raises an OSError instead of returning empty bytes
+            except OSError as err:
+                raise OSError("KMS connection closed") from err
             except BLOCKING_IO_ERRORS:
                 raise socket.timeout("timed out") from None
             finally:

pymongo/network_layer.py

@@ -130,7 +130,7 @@ def _is_ready(fut: Future) -> None:
                         loop.remove_writer(fd)
 
     async def _async_receive_ssl(
-        conn: _sslConn, length: int, loop: AbstractEventLoop
+        conn: _sslConn, length: int, loop: AbstractEventLoop, once: Optional[bool] = False
     ) -> memoryview:
         mv = memoryview(bytearray(length))
         total_read = 0
@@ -145,6 +145,9 @@ def _is_ready(fut: Future) -> None:
                 read = conn.recv_into(mv[total_read:])
                 if read == 0:
                     raise OSError("connection closed")
+                # KMS responses update their expected size after the first batch, stop reading after one loop
+                if once:
+                    return mv[:read]
                 total_read += read
             except BLOCKING_IO_ERRORS as exc:
                 fd = conn.fileno()
@@ -275,6 +278,28 @@ async def async_receive_data(
         sock.settimeout(sock_timeout)
 
 
+async def async_receive_data_socket(
+    sock: Union[socket.socket, _sslConn], length: int
+) -> memoryview:
+    sock_timeout = sock.gettimeout()
+    timeout = sock_timeout
+
+    sock.settimeout(0.0)
+    loop = asyncio.get_event_loop()
+    try:
+        if _HAVE_SSL and isinstance(sock, (SSLSocket, _sslConn)):
+            return await asyncio.wait_for(
+                _async_receive_ssl(sock, length, loop, once=True),  # type: ignore[arg-type]
+                timeout=timeout,
+            )
+        else:
+            return await asyncio.wait_for(_async_receive(sock, length, loop), timeout=timeout)  # type: ignore[arg-type]
+    except asyncio.TimeoutError as err:
+        raise socket.timeout("timed out") from err
+    finally:
+        sock.settimeout(sock_timeout)
+
+
 async def _async_receive(conn: socket.socket, length: int, loop: AbstractEventLoop) -> memoryview:
     mv = memoryview(bytearray(length))
     bytes_read = 0

pymongo/synchronous/encryption.py

@@ -180,10 +180,20 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                 while kms_context.bytes_needed > 0:
                     # CSOT: update timeout.
                     conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
-                    data = conn.recv(kms_context.bytes_needed)
+                    if _IS_SYNC:
+                        data = conn.recv(kms_context.bytes_needed)
+                    else:
+                        from pymongo.network_layer import (  # type: ignore[attr-defined]
+                            receive_data_socket,
+                        )
+
+                        data = receive_data_socket(conn, kms_context.bytes_needed)
                     if not data:
                         raise OSError("KMS connection closed")
                     kms_context.feed(data)
+            # Async raises an OSError instead of returning empty bytes
+            except OSError as err:
+                raise OSError("KMS connection closed") from err
             except BLOCKING_IO_ERRORS:
                 raise socket.timeout("timed out") from None
             finally:

test/__init__.py

@@ -464,11 +464,12 @@ def wrap(*args, **kwargs):
                 if not self.connected:
                     pair = self.pair
                     raise SkipTest(f"Cannot connect to MongoDB on {pair}")
-                if iscoroutinefunction(condition) and condition():
-                    if wraps_async:
-                        return f(*args, **kwargs)
-                    else:
-                        return f(*args, **kwargs)
+                if iscoroutinefunction(condition):
+                    if condition():
+                        if wraps_async:
+                            return f(*args, **kwargs)
+                        else:
+                            return f(*args, **kwargs)
                 elif condition():
                     if wraps_async:
                         return f(*args, **kwargs)

test/asynchronous/__init__.py

@@ -466,11 +466,12 @@ async def wrap(*args, **kwargs):
                 if not self.connected:
                     pair = await self.pair
                     raise SkipTest(f"Cannot connect to MongoDB on {pair}")
-                if iscoroutinefunction(condition) and await condition():
-                    if wraps_async:
-                        return await f(*args, **kwargs)
-                    else:
-                        return f(*args, **kwargs)
+                if iscoroutinefunction(condition):
+                    if await condition():
+                        if wraps_async:
+                            return await f(*args, **kwargs)
+                        else:
+                            return f(*args, **kwargs)
                 elif condition():
                     if wraps_async:
                         return await f(*args, **kwargs)