pytest_plugin.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191
  1. from __future__ import annotations
  2. import sys
  3. from collections.abc import Generator, Iterator
  4. from contextlib import ExitStack, contextmanager
  5. from inspect import isasyncgenfunction, iscoroutinefunction, ismethod
  6. from typing import Any, Dict, Tuple, cast
  7. import pytest
  8. import sniffio
  9. from _pytest.fixtures import SubRequest
  10. from _pytest.outcomes import Exit
  11. from ._core._eventloop import get_all_backends, get_async_backend
  12. from ._core._exceptions import iterate_exceptions
  13. from .abc import TestRunner
  14. if sys.version_info < (3, 11):
  15. from exceptiongroup import ExceptionGroup
  16. _current_runner: TestRunner | None = None
  17. _runner_stack: ExitStack | None = None
  18. _runner_leases = 0
  19. def extract_backend_and_options(backend: object) -> tuple[str, dict[str, Any]]:
  20. if isinstance(backend, str):
  21. return backend, {}
  22. elif isinstance(backend, tuple) and len(backend) == 2:
  23. if isinstance(backend[0], str) and isinstance(backend[1], dict):
  24. return cast(Tuple[str, Dict[str, Any]], backend)
  25. raise TypeError("anyio_backend must be either a string or tuple of (string, dict)")
  26. @contextmanager
  27. def get_runner(
  28. backend_name: str, backend_options: dict[str, Any]
  29. ) -> Iterator[TestRunner]:
  30. global _current_runner, _runner_leases, _runner_stack
  31. if _current_runner is None:
  32. asynclib = get_async_backend(backend_name)
  33. _runner_stack = ExitStack()
  34. if sniffio.current_async_library_cvar.get(None) is None:
  35. # Since we're in control of the event loop, we can cache the name of the
  36. # async library
  37. token = sniffio.current_async_library_cvar.set(backend_name)
  38. _runner_stack.callback(sniffio.current_async_library_cvar.reset, token)
  39. backend_options = backend_options or {}
  40. _current_runner = _runner_stack.enter_context(
  41. asynclib.create_test_runner(backend_options)
  42. )
  43. _runner_leases += 1
  44. try:
  45. yield _current_runner
  46. finally:
  47. _runner_leases -= 1
  48. if not _runner_leases:
  49. assert _runner_stack is not None
  50. _runner_stack.close()
  51. _runner_stack = _current_runner = None
  52. def pytest_configure(config: Any) -> None:
  53. config.addinivalue_line(
  54. "markers",
  55. "anyio: mark the (coroutine function) test to be run "
  56. "asynchronously via anyio.",
  57. )
  58. @pytest.hookimpl(hookwrapper=True)
  59. def pytest_fixture_setup(fixturedef: Any, request: Any) -> Generator[Any]:
  60. def wrapper(
  61. *args: Any, anyio_backend: Any, request: SubRequest, **kwargs: Any
  62. ) -> Any:
  63. # Rebind any fixture methods to the request instance
  64. if (
  65. request.instance
  66. and ismethod(func)
  67. and type(func.__self__) is type(request.instance)
  68. ):
  69. local_func = func.__func__.__get__(request.instance)
  70. else:
  71. local_func = func
  72. backend_name, backend_options = extract_backend_and_options(anyio_backend)
  73. if has_backend_arg:
  74. kwargs["anyio_backend"] = anyio_backend
  75. if has_request_arg:
  76. kwargs["request"] = request
  77. with get_runner(backend_name, backend_options) as runner:
  78. if isasyncgenfunction(local_func):
  79. yield from runner.run_asyncgen_fixture(local_func, kwargs)
  80. else:
  81. yield runner.run_fixture(local_func, kwargs)
  82. # Only apply this to coroutine functions and async generator functions in requests
  83. # that involve the anyio_backend fixture
  84. func = fixturedef.func
  85. if isasyncgenfunction(func) or iscoroutinefunction(func):
  86. if "anyio_backend" in request.fixturenames:
  87. fixturedef.func = wrapper
  88. original_argname = fixturedef.argnames
  89. if not (has_backend_arg := "anyio_backend" in fixturedef.argnames):
  90. fixturedef.argnames += ("anyio_backend",)
  91. if not (has_request_arg := "request" in fixturedef.argnames):
  92. fixturedef.argnames += ("request",)
  93. try:
  94. return (yield)
  95. finally:
  96. fixturedef.func = func
  97. fixturedef.argnames = original_argname
  98. return (yield)
  99. @pytest.hookimpl(tryfirst=True)
  100. def pytest_pycollect_makeitem(collector: Any, name: Any, obj: Any) -> None:
  101. if collector.istestfunction(obj, name):
  102. inner_func = obj.hypothesis.inner_test if hasattr(obj, "hypothesis") else obj
  103. if iscoroutinefunction(inner_func):
  104. marker = collector.get_closest_marker("anyio")
  105. own_markers = getattr(obj, "pytestmark", ())
  106. if marker or any(marker.name == "anyio" for marker in own_markers):
  107. pytest.mark.usefixtures("anyio_backend")(obj)
  108. @pytest.hookimpl(tryfirst=True)
  109. def pytest_pyfunc_call(pyfuncitem: Any) -> bool | None:
  110. def run_with_hypothesis(**kwargs: Any) -> None:
  111. with get_runner(backend_name, backend_options) as runner:
  112. runner.run_test(original_func, kwargs)
  113. backend = pyfuncitem.funcargs.get("anyio_backend")
  114. if backend:
  115. backend_name, backend_options = extract_backend_and_options(backend)
  116. if hasattr(pyfuncitem.obj, "hypothesis"):
  117. # Wrap the inner test function unless it's already wrapped
  118. original_func = pyfuncitem.obj.hypothesis.inner_test
  119. if original_func.__qualname__ != run_with_hypothesis.__qualname__:
  120. if iscoroutinefunction(original_func):
  121. pyfuncitem.obj.hypothesis.inner_test = run_with_hypothesis
  122. return None
  123. if iscoroutinefunction(pyfuncitem.obj):
  124. funcargs = pyfuncitem.funcargs
  125. testargs = {arg: funcargs[arg] for arg in pyfuncitem._fixtureinfo.argnames}
  126. with get_runner(backend_name, backend_options) as runner:
  127. try:
  128. runner.run_test(pyfuncitem.obj, testargs)
  129. except ExceptionGroup as excgrp:
  130. for exc in iterate_exceptions(excgrp):
  131. if isinstance(exc, (Exit, KeyboardInterrupt, SystemExit)):
  132. raise exc from excgrp
  133. raise
  134. return True
  135. return None
  136. @pytest.fixture(scope="module", params=get_all_backends())
  137. def anyio_backend(request: Any) -> Any:
  138. return request.param
  139. @pytest.fixture
  140. def anyio_backend_name(anyio_backend: Any) -> str:
  141. if isinstance(anyio_backend, str):
  142. return anyio_backend
  143. else:
  144. return anyio_backend[0]
  145. @pytest.fixture
  146. def anyio_backend_options(anyio_backend: Any) -> dict[str, Any]:
  147. if isinstance(anyio_backend, str):
  148. return {}
  149. else:
  150. return anyio_backend[1]