codekingpro/portable-devtools

1# testing/asyncio.py
2# Copyright (C) 2005-2024 the SQLAlchemy authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of SQLAlchemy and is released under
6# the MIT License: https://www.opensource.org/licenses/mit-license.php
7# mypy: ignore-errors
8
9
10# functions and wrappers to run tests, fixtures, provisioning and
11# setup/teardown in an asyncio event loop, conditionally based on the
12# current DB driver being used for a test.
13
14# note that SQLAlchemy's asyncio integration also supports a method
15# of running individual asyncio functions inside of separate event loops
16# using "async_fallback" mode; however running whole functions in the event
17# loop is a more accurate test for how SQLAlchemy's asyncio features
18# would run in the real world.
19
20
21from __future__ import annotations
22
23from functools import wraps
24import inspect
25
26from . import config
27from ..util.concurrency import _AsyncUtil
28
29# may be set to False if the
30# --disable-asyncio flag is passed to the test runner.
31ENABLE_ASYNCIO = True
32_async_util = _AsyncUtil()  # it has lazy init so just always create one
33
34
35def _shutdown():
36    """called when the test finishes"""
37    _async_util.close()
38
39
40def _run_coroutine_function(fn, *args, **kwargs):
41    return _async_util.run(fn, *args, **kwargs)
42
43
44def _assume_async(fn, *args, **kwargs):
45    """Run a function in an asyncio loop unconditionally.
46
47    This function is used for provisioning features like
48    testing a database connection for server info.
49
50    Note that for blocking IO database drivers, this means they block the
51    event loop.
52
53    """
54
55    if not ENABLE_ASYNCIO:
56        return fn(*args, **kwargs)
57
58    return _async_util.run_in_greenlet(fn, *args, **kwargs)
59
60
61def _maybe_async_provisioning(fn, *args, **kwargs):
62    """Run a function in an asyncio loop if any current drivers might need it.
63
64    This function is used for provisioning features that take
65    place outside of a specific database driver being selected, so if the
66    current driver that happens to be used for the provisioning operation
67    is an async driver, it will run in asyncio and not fail.
68
69    Note that for blocking IO database drivers, this means they block the
70    event loop.
71
72    """
73    if not ENABLE_ASYNCIO:
74        return fn(*args, **kwargs)
75
76    if config.any_async:
77        return _async_util.run_in_greenlet(fn, *args, **kwargs)
78    else:
79        return fn(*args, **kwargs)
80
81
82def _maybe_async(fn, *args, **kwargs):
83    """Run a function in an asyncio loop if the current selected driver is
84    async.
85
86    This function is used for test setup/teardown and tests themselves
87    where the current DB driver is known.
88
89
90    """
91    if not ENABLE_ASYNCIO:
92        return fn(*args, **kwargs)
93
94    is_async = config._current.is_async
95
96    if is_async:
97        return _async_util.run_in_greenlet(fn, *args, **kwargs)
98    else:
99        return fn(*args, **kwargs)
100
101
102def _maybe_async_wrapper(fn):
103    """Apply the _maybe_async function to an existing function and return
104    as a wrapped callable, supporting generator functions as well.
105
106    This is currently used for pytest fixtures that support generator use.
107
108    """
109
110    if inspect.isgeneratorfunction(fn):
111        _stop = object()
112
113        def call_next(gen):
114            try:
115                return next(gen)
116                # can't raise StopIteration in an awaitable.
117            except StopIteration:
118                return _stop
119
120        @wraps(fn)
121        def wrap_fixture(*args, **kwargs):
122            gen = fn(*args, **kwargs)
123            while True:
124                value = _maybe_async(call_next, gen)
125                if value is _stop:
126                    break
127                yield value
128
129    else:
130
131        @wraps(fn)
132        def wrap_fixture(*args, **kwargs):
133            return _maybe_async(fn, *args, **kwargs)
134
135    return wrap_fixture
136