Update 2025-04-13_16:25:39

venv/lib/python3.11/site-packages/sqlalchemy/testing/asyncio.py

@@ -0,0 +1,135 @@
+# testing/asyncio.py
+# Copyright (C) 2005-2025 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+# mypy: ignore-errors
+
+
+# functions and wrappers to run tests, fixtures, provisioning and
+# setup/teardown in an asyncio event loop, conditionally based on the
+# current DB driver being used for a test.
+
+# note that SQLAlchemy's asyncio integration also supports a method
+# of running individual asyncio functions inside of separate event loops
+# using "async_fallback" mode; however running whole functions in the event
+# loop is a more accurate test for how SQLAlchemy's asyncio features
+# would run in the real world.
+
+
+from __future__ import annotations
+
+from functools import wraps
+import inspect
+
+from . import config
+from ..util.concurrency import _AsyncUtil
+
+# may be set to False if the
+# --disable-asyncio flag is passed to the test runner.
+ENABLE_ASYNCIO = True
+_async_util = _AsyncUtil()  # it has lazy init so just always create one
+
+
+def _shutdown():
+    """called when the test finishes"""
+    _async_util.close()
+
+
+def _run_coroutine_function(fn, *args, **kwargs):
+    return _async_util.run(fn, *args, **kwargs)
+
+
+def _assume_async(fn, *args, **kwargs):
+    """Run a function in an asyncio loop unconditionally.
+
+    This function is used for provisioning features like
+    testing a database connection for server info.
+
+    Note that for blocking IO database drivers, this means they block the
+    event loop.
+
+    """
+
+    if not ENABLE_ASYNCIO:
+        return fn(*args, **kwargs)
+
+    return _async_util.run_in_greenlet(fn, *args, **kwargs)
+
+
+def _maybe_async_provisioning(fn, *args, **kwargs):
+    """Run a function in an asyncio loop if any current drivers might need it.
+
+    This function is used for provisioning features that take
+    place outside of a specific database driver being selected, so if the
+    current driver that happens to be used for the provisioning operation
+    is an async driver, it will run in asyncio and not fail.
+
+    Note that for blocking IO database drivers, this means they block the
+    event loop.
+
+    """
+    if not ENABLE_ASYNCIO:
+        return fn(*args, **kwargs)
+
+    if config.any_async:
+        return _async_util.run_in_greenlet(fn, *args, **kwargs)
+    else:
+        return fn(*args, **kwargs)
+
+
+def _maybe_async(fn, *args, **kwargs):
+    """Run a function in an asyncio loop if the current selected driver is
+    async.
+
+    This function is used for test setup/teardown and tests themselves
+    where the current DB driver is known.
+
+
+    """
+    if not ENABLE_ASYNCIO:
+        return fn(*args, **kwargs)
+
+    is_async = config._current.is_async
+
+    if is_async:
+        return _async_util.run_in_greenlet(fn, *args, **kwargs)
+    else:
+        return fn(*args, **kwargs)
+
+
+def _maybe_async_wrapper(fn):
+    """Apply the _maybe_async function to an existing function and return
+    as a wrapped callable, supporting generator functions as well.
+
+    This is currently used for pytest fixtures that support generator use.
+
+    """
+
+    if inspect.isgeneratorfunction(fn):
+        _stop = object()
+
+        def call_next(gen):
+            try:
+                return next(gen)
+                # can't raise StopIteration in an awaitable.
+            except StopIteration:
+                return _stop
+
+        @wraps(fn)
+        def wrap_fixture(*args, **kwargs):
+            gen = fn(*args, **kwargs)
+            while True:
+                value = _maybe_async(call_next, gen)
+                if value is _stop:
+                    break
+                yield value
+
+    else:
+
+        @wraps(fn)
+        def wrap_fixture(*args, **kwargs):
+            return _maybe_async(fn, *args, **kwargs)
+
+    return wrap_fixture