codekingpro/portable-devtools

1from __future__ import annotations2 3import dataclasses4import socket5import sys6from collections.abc import Callable, Generator, Iterator7from contextlib import ExitStack, contextmanager8from inspect import isasyncgenfunction, iscoroutinefunction, ismethod9from typing import Any, cast10 11import pytest12from _pytest.fixtures import FuncFixtureInfo, SubRequest13from _pytest.outcomes import Exit14from _pytest.python import CallSpec215from _pytest.scope import Scope16 17from . import get_available_backends18from ._core._eventloop import (19    current_async_library,20    get_async_backend,21    reset_current_async_library,22    set_current_async_library,23)24from ._core._exceptions import iterate_exceptions25from .abc import TestRunner26 27if sys.version_info < (3, 11):28    from exceptiongroup import ExceptionGroup29 30_current_runner: TestRunner | None = None31_runner_stack: ExitStack | None = None32_runner_leases = 033 34 35def extract_backend_and_options(backend: object) -> tuple[str, dict[str, Any]]:36    if isinstance(backend, str):37        return backend, {}38    elif isinstance(backend, tuple) and len(backend) == 2:39        if isinstance(backend[0], str) and isinstance(backend[1], dict):40            return cast(tuple[str, dict[str, Any]], backend)41 42    raise TypeError("anyio_backend must be either a string or tuple of (string, dict)")43 44 45@contextmanager46def get_runner(47    backend_name: str, backend_options: dict[str, Any]48) -> Iterator[TestRunner]:49    global _current_runner, _runner_leases, _runner_stack50    if _current_runner is None:51        asynclib = get_async_backend(backend_name)52        _runner_stack = ExitStack()53        if current_async_library() is None:54            # Since we're in control of the event loop, we can cache the name of the55            # async library56            token = set_current_async_library(backend_name)57            _runner_stack.callback(reset_current_async_library, token)58 59        backend_options = backend_options or {}60        _current_runner = _runner_stack.enter_context(61            asynclib.create_test_runner(backend_options)62        )63 64    _runner_leases += 165    try:66        yield _current_runner67    finally:68        _runner_leases -= 169        if not _runner_leases:70            assert _runner_stack is not None71            _runner_stack.close()72            _runner_stack = _current_runner = None73 74 75def pytest_addoption(parser: pytest.Parser) -> None:76    parser.addini(77        "anyio_mode",78        default="strict",79        help='AnyIO plugin mode (either "strict" or "auto")',80    )81 82 83def pytest_configure(config: pytest.Config) -> None:84    config.addinivalue_line(85        "markers",86        "anyio: mark the (coroutine function) test to be run asynchronously via anyio.",87    )88    if (89        config.getini("anyio_mode") == "auto"90        and config.pluginmanager.has_plugin("asyncio")91        and config.getini("asyncio_mode") == "auto"92    ):93        config.issue_config_time_warning(94            pytest.PytestConfigWarning(95                "AnyIO auto mode has been enabled together with pytest-asyncio auto "96                "mode. This may cause unexpected behavior."97            ),98            1,99        )100 101 102@pytest.hookimpl(hookwrapper=True)103def pytest_fixture_setup(fixturedef: Any, request: Any) -> Generator[Any]:104    def wrapper(anyio_backend: Any, request: SubRequest, **kwargs: Any) -> Any:105        # Rebind any fixture methods to the request instance106        if (107            request.instance108            and ismethod(func)109            and type(func.__self__) is type(request.instance)110        ):111            local_func = func.__func__.__get__(request.instance)112        else:113            local_func = func114 115        backend_name, backend_options = extract_backend_and_options(anyio_backend)116        if has_backend_arg:117            kwargs["anyio_backend"] = anyio_backend118 119        if has_request_arg:120            kwargs["request"] = request121 122        with get_runner(backend_name, backend_options) as runner:123            if isasyncgenfunction(local_func):124                yield from runner.run_asyncgen_fixture(local_func, kwargs)125            else:126                yield runner.run_fixture(local_func, kwargs)127 128    # Only apply this to coroutine functions and async generator functions in requests129    # that involve the anyio_backend fixture130    func = fixturedef.func131    if isasyncgenfunction(func) or iscoroutinefunction(func):132        if "anyio_backend" in request.fixturenames:133            fixturedef.func = wrapper134            original_argname = fixturedef.argnames135 136            if not (has_backend_arg := "anyio_backend" in fixturedef.argnames):137                fixturedef.argnames += ("anyio_backend",)138 139            if not (has_request_arg := "request" in fixturedef.argnames):140                fixturedef.argnames += ("request",)141 142            try:143                return (yield)144            finally:145                fixturedef.func = func146                fixturedef.argnames = original_argname147 148    return (yield)149 150 151@pytest.hookimpl(tryfirst=True)152def pytest_pycollect_makeitem(153    collector: pytest.Module | pytest.Class, name: str, obj: object154) -> None:155    if collector.istestfunction(obj, name):156        inner_func = obj.hypothesis.inner_test if hasattr(obj, "hypothesis") else obj157        if iscoroutinefunction(inner_func):158            anyio_auto_mode = collector.config.getini("anyio_mode") == "auto"159            marker = collector.get_closest_marker("anyio")160            own_markers = getattr(obj, "pytestmark", ())161            if (162                anyio_auto_mode163                or marker164                or any(marker.name == "anyio" for marker in own_markers)165            ):166                pytest.mark.usefixtures("anyio_backend")(obj)167 168 169def pytest_collection_finish(session: pytest.Session) -> None:170    for i, item in reversed(list(enumerate(session.items))):171        if (172            isinstance(item, pytest.Function)173            and iscoroutinefunction(item.function)174            and item.get_closest_marker("anyio") is not None175            and "anyio_backend" not in item.fixturenames176        ):177            new_items = []178            try:179                cs_fields = {f.name for f in dataclasses.fields(CallSpec2)}180            except TypeError:181                cs_fields = set()182 183            for param_index, backend in enumerate(get_available_backends()):184                if "_arg2scope" in cs_fields:  # pytest >= 8185                    callspec = CallSpec2(186                        params={"anyio_backend": backend},187                        indices={"anyio_backend": param_index},188                        _arg2scope={"anyio_backend": Scope.Module},189                        _idlist=[backend],190                        marks=[],191                    )192                else:  # pytest 7.x193                    callspec = CallSpec2(  # type: ignore[call-arg]194                        funcargs={},195                        params={"anyio_backend": backend},196                        indices={"anyio_backend": param_index},197                        arg2scope={"anyio_backend": Scope.Module},198                        idlist=[backend],199                        marks=[],200                    )201 202                fi = item._fixtureinfo203                new_names_closure = list(fi.names_closure)204                if "anyio_backend" not in new_names_closure:205                    new_names_closure.append("anyio_backend")206 207                new_fixtureinfo = FuncFixtureInfo(208                    argnames=fi.argnames,209                    initialnames=fi.initialnames,210                    names_closure=new_names_closure,211                    name2fixturedefs=fi.name2fixturedefs,212                )213                new_item = pytest.Function.from_parent(214                    item.parent,215                    name=f"{item.originalname}[{backend}]",216                    callspec=callspec,217                    callobj=item.obj,218                    fixtureinfo=new_fixtureinfo,219                    keywords=item.keywords,220                    originalname=item.originalname,221                )222                new_items.append(new_item)223 224            session.items[i : i + 1] = new_items225 226 227@pytest.hookimpl(tryfirst=True)228def pytest_pyfunc_call(pyfuncitem: Any) -> bool | None:229    def run_with_hypothesis(**kwargs: Any) -> None:230        with get_runner(backend_name, backend_options) as runner:231            runner.run_test(original_func, kwargs)232 233    backend = pyfuncitem.funcargs.get("anyio_backend")234    if backend:235        backend_name, backend_options = extract_backend_and_options(backend)236 237        if hasattr(pyfuncitem.obj, "hypothesis"):238            # Wrap the inner test function unless it's already wrapped239            original_func = pyfuncitem.obj.hypothesis.inner_test240            if original_func.__qualname__ != run_with_hypothesis.__qualname__:241                if iscoroutinefunction(original_func):242                    pyfuncitem.obj.hypothesis.inner_test = run_with_hypothesis243 244            return None245 246        if iscoroutinefunction(pyfuncitem.obj):247            funcargs = pyfuncitem.funcargs248            testargs = {arg: funcargs[arg] for arg in pyfuncitem._fixtureinfo.argnames}249            with get_runner(backend_name, backend_options) as runner:250                try:251                    runner.run_test(pyfuncitem.obj, testargs)252                except ExceptionGroup as excgrp:253                    for exc in iterate_exceptions(excgrp):254                        if isinstance(exc, (Exit, KeyboardInterrupt, SystemExit)):255                            raise exc from excgrp256 257                    raise258 259            return True260 261    return None262 263 264@pytest.fixture(scope="module", params=get_available_backends())265def anyio_backend(request: Any) -> Any:266    return request.param267 268 269@pytest.fixture270def anyio_backend_name(anyio_backend: Any) -> str:271    if isinstance(anyio_backend, str):272        return anyio_backend273    else:274        return anyio_backend[0]275 276 277@pytest.fixture278def anyio_backend_options(anyio_backend: Any) -> dict[str, Any]:279    if isinstance(anyio_backend, str):280        return {}281    else:282        return anyio_backend[1]283 284 285class FreePortFactory:286    """287    Manages port generation based on specified socket kind, ensuring no duplicate288    ports are generated.289 290    This class provides functionality for generating available free ports on the291    system. It is initialized with a specific socket kind and can generate ports292    for given address families while avoiding reuse of previously generated ports.293 294    Users should not instantiate this class directly, but use the295    ``free_tcp_port_factory`` and ``free_udp_port_factory`` fixtures instead. For simple296    uses cases, ``free_tcp_port`` and ``free_udp_port`` can be used instead.297    """298 299    def __init__(self, kind: socket.SocketKind) -> None:300        self._kind = kind301        self._generated = set[int]()302 303    @property304    def kind(self) -> socket.SocketKind:305        """306        The type of socket connection (e.g., :data:`~socket.SOCK_STREAM` or307        :data:`~socket.SOCK_DGRAM`) used to bind for checking port availability308 309        """310        return self._kind311 312    def __call__(self, family: socket.AddressFamily | None = None) -> int:313        """314        Return an unbound port for the given address family.315 316        :param family: if omitted, both IPv4 and IPv6 addresses will be tried317        :return: a port number318 319        """320        if family is not None:321            families = [family]322        else:323            families = [socket.AF_INET]324            if socket.has_ipv6:325                families.append(socket.AF_INET6)326 327        while True:328            port = 0329            with ExitStack() as stack:330                for family in families:331                    sock = stack.enter_context(socket.socket(family, self._kind))332                    addr = "::1" if family == socket.AF_INET6 else "127.0.0.1"333                    try:334                        sock.bind((addr, port))335                    except OSError:336                        break337 338                    if not port:339                        port = sock.getsockname()[1]340                else:341                    if port not in self._generated:342                        self._generated.add(port)343                        return port344 345 346@pytest.fixture(scope="session")347def free_tcp_port_factory() -> FreePortFactory:348    return FreePortFactory(socket.SOCK_STREAM)349 350 351@pytest.fixture(scope="session")352def free_udp_port_factory() -> FreePortFactory:353    return FreePortFactory(socket.SOCK_DGRAM)354 355 356@pytest.fixture357def free_tcp_port(free_tcp_port_factory: Callable[[], int]) -> int:358    return free_tcp_port_factory()359 360 361@pytest.fixture362def free_udp_port(free_udp_port_factory: Callable[[], int]) -> int:363    return free_udp_port_factory()364