asyncio.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128
  1. # testing/asyncio.py
  2. # Copyright (C) 2005-2022 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. # functions and wrappers to run tests, fixtures, provisioning and
  8. # setup/teardown in an asyncio event loop, conditionally based on the
  9. # current DB driver being used for a test.
  10. # note that SQLAlchemy's asyncio integration also supports a method
  11. # of running individual asyncio functions inside of separate event loops
  12. # using "async_fallback" mode; however running whole functions in the event
  13. # loop is a more accurate test for how SQLAlchemy's asyncio features
  14. # would run in the real world.
  15. from functools import wraps
  16. import inspect
  17. from . import config
  18. from ..util.concurrency import _util_async_run
  19. from ..util.concurrency import _util_async_run_coroutine_function
  20. # may be set to False if the
  21. # --disable-asyncio flag is passed to the test runner.
  22. ENABLE_ASYNCIO = True
  23. def _run_coroutine_function(fn, *args, **kwargs):
  24. return _util_async_run_coroutine_function(fn, *args, **kwargs)
  25. def _assume_async(fn, *args, **kwargs):
  26. """Run a function in an asyncio loop unconditionally.
  27. This function is used for provisioning features like
  28. testing a database connection for server info.
  29. Note that for blocking IO database drivers, this means they block the
  30. event loop.
  31. """
  32. if not ENABLE_ASYNCIO:
  33. return fn(*args, **kwargs)
  34. return _util_async_run(fn, *args, **kwargs)
  35. def _maybe_async_provisioning(fn, *args, **kwargs):
  36. """Run a function in an asyncio loop if any current drivers might need it.
  37. This function is used for provisioning features that take
  38. place outside of a specific database driver being selected, so if the
  39. current driver that happens to be used for the provisioning operation
  40. is an async driver, it will run in asyncio and not fail.
  41. Note that for blocking IO database drivers, this means they block the
  42. event loop.
  43. """
  44. if not ENABLE_ASYNCIO:
  45. return fn(*args, **kwargs)
  46. if config.any_async:
  47. return _util_async_run(fn, *args, **kwargs)
  48. else:
  49. return fn(*args, **kwargs)
  50. def _maybe_async(fn, *args, **kwargs):
  51. """Run a function in an asyncio loop if the current selected driver is
  52. async.
  53. This function is used for test setup/teardown and tests themselves
  54. where the current DB driver is known.
  55. """
  56. if not ENABLE_ASYNCIO:
  57. return fn(*args, **kwargs)
  58. is_async = config._current.is_async
  59. if is_async:
  60. return _util_async_run(fn, *args, **kwargs)
  61. else:
  62. return fn(*args, **kwargs)
  63. def _maybe_async_wrapper(fn):
  64. """Apply the _maybe_async function to an existing function and return
  65. as a wrapped callable, supporting generator functions as well.
  66. This is currently used for pytest fixtures that support generator use.
  67. """
  68. if inspect.isgeneratorfunction(fn):
  69. _stop = object()
  70. def call_next(gen):
  71. try:
  72. return next(gen)
  73. # can't raise StopIteration in an awaitable.
  74. except StopIteration:
  75. return _stop
  76. @wraps(fn)
  77. def wrap_fixture(*args, **kwargs):
  78. gen = fn(*args, **kwargs)
  79. while True:
  80. value = _maybe_async(call_next, gen)
  81. if value is _stop:
  82. break
  83. yield value
  84. else:
  85. @wraps(fn)
  86. def wrap_fixture(*args, **kwargs):
  87. return _maybe_async(fn, *args, **kwargs)
  88. return wrap_fixture