codekingpro/portable-devtools

1from __future__ import annotations2 3import importlib.metadata4import typing as t5from contextlib import contextmanager6from contextlib import ExitStack7from copy import copy8from types import TracebackType9from urllib.parse import urlsplit10 11import werkzeug.test12from click.testing import CliRunner13from click.testing import Result14from werkzeug.test import Client15from werkzeug.wrappers import Request as BaseRequest16 17from .cli import ScriptInfo18from .sessions import SessionMixin19 20if t.TYPE_CHECKING:  # pragma: no cover21    from _typeshed.wsgi import WSGIEnvironment22    from werkzeug.test import TestResponse23 24    from .app import Flask25 26 27class EnvironBuilder(werkzeug.test.EnvironBuilder):28    """An :class:`~werkzeug.test.EnvironBuilder`, that takes defaults from the29    application.30 31    :param app: The Flask application to configure the environment from.32    :param path: URL path being requested.33    :param base_url: Base URL where the app is being served, which34        ``path`` is relative to. If not given, built from35        :data:`PREFERRED_URL_SCHEME`, ``subdomain``,36        :data:`SERVER_NAME`, and :data:`APPLICATION_ROOT`.37    :param subdomain: Subdomain name to append to :data:`SERVER_NAME`.38    :param url_scheme: Scheme to use instead of39        :data:`PREFERRED_URL_SCHEME`.40    :param json: If given, this is serialized as JSON and passed as41        ``data``. Also defaults ``content_type`` to42        ``application/json``.43    :param args: other positional arguments passed to44        :class:`~werkzeug.test.EnvironBuilder`.45    :param kwargs: other keyword arguments passed to46        :class:`~werkzeug.test.EnvironBuilder`.47    """48 49    def __init__(50        self,51        app: Flask,52        path: str = "/",53        base_url: str | None = None,54        subdomain: str | None = None,55        url_scheme: str | None = None,56        *args: t.Any,57        **kwargs: t.Any,58    ) -> None:59        assert not (base_url or subdomain or url_scheme) or (60            base_url is not None61        ) != bool(subdomain or url_scheme), (62            'Cannot pass "subdomain" or "url_scheme" with "base_url".'63        )64 65        if base_url is None:66            http_host = app.config.get("SERVER_NAME") or "localhost"67            app_root = app.config["APPLICATION_ROOT"]68 69            if subdomain:70                http_host = f"{subdomain}.{http_host}"71 72            if url_scheme is None:73                url_scheme = app.config["PREFERRED_URL_SCHEME"]74 75            url = urlsplit(path)76            base_url = (77                f"{url.scheme or url_scheme}://{url.netloc or http_host}"78                f"/{app_root.lstrip('/')}"79            )80            path = url.path81 82            if url.query:83                path = f"{path}?{url.query}"84 85        self.app = app86        super().__init__(path, base_url, *args, **kwargs)87 88    def json_dumps(self, obj: t.Any, **kwargs: t.Any) -> str:89        """Serialize ``obj`` to a JSON-formatted string.90 91        The serialization will be configured according to the config associated92        with this EnvironBuilder's ``app``.93        """94        return self.app.json.dumps(obj, **kwargs)95 96 97_werkzeug_version = ""98 99 100def _get_werkzeug_version() -> str:101    global _werkzeug_version102 103    if not _werkzeug_version:104        _werkzeug_version = importlib.metadata.version("werkzeug")105 106    return _werkzeug_version107 108 109class FlaskClient(Client):110    """Works like a regular Werkzeug test client but has knowledge about111    Flask's contexts to defer the cleanup of the request context until112    the end of a ``with`` block. For general information about how to113    use this class refer to :class:`werkzeug.test.Client`.114 115    .. versionchanged:: 0.12116       `app.test_client()` includes preset default environment, which can be117       set after instantiation of the `app.test_client()` object in118       `client.environ_base`.119 120    Basic usage is outlined in the :doc:`/testing` chapter.121    """122 123    application: Flask124 125    def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:126        super().__init__(*args, **kwargs)127        self.preserve_context = False128        self._new_contexts: list[t.ContextManager[t.Any]] = []129        self._context_stack = ExitStack()130        self.environ_base = {131            "REMOTE_ADDR": "127.0.0.1",132            "HTTP_USER_AGENT": f"Werkzeug/{_get_werkzeug_version()}",133        }134 135    @contextmanager136    def session_transaction(137        self, *args: t.Any, **kwargs: t.Any138    ) -> t.Iterator[SessionMixin]:139        """When used in combination with a ``with`` statement this opens a140        session transaction.  This can be used to modify the session that141        the test client uses.  Once the ``with`` block is left the session is142        stored back.143 144        ::145 146            with client.session_transaction() as session:147                session['value'] = 42148 149        Internally this is implemented by going through a temporary test150        request context and since session handling could depend on151        request variables this function accepts the same arguments as152        :meth:`~flask.Flask.test_request_context` which are directly153        passed through.154        """155        if self._cookies is None:156            raise TypeError(157                "Cookies are disabled. Create a client with 'use_cookies=True'."158            )159 160        app = self.application161        ctx = app.test_request_context(*args, **kwargs)162        self._add_cookies_to_wsgi(ctx.request.environ)163 164        with ctx:165            sess = app.session_interface.open_session(app, ctx.request)166 167        if sess is None:168            raise RuntimeError("Session backend did not open a session.")169 170        yield sess171        resp = app.response_class()172 173        if app.session_interface.is_null_session(sess):174            return175 176        with ctx:177            app.session_interface.save_session(app, sess, resp)178 179        self._update_cookies_from_response(180            ctx.request.host.partition(":")[0],181            ctx.request.path,182            resp.headers.getlist("Set-Cookie"),183        )184 185    def _copy_environ(self, other: WSGIEnvironment) -> WSGIEnvironment:186        out = {**self.environ_base, **other}187 188        if self.preserve_context:189            out["werkzeug.debug.preserve_context"] = self._new_contexts.append190 191        return out192 193    def _request_from_builder_args(194        self, args: tuple[t.Any, ...], kwargs: dict[str, t.Any]195    ) -> BaseRequest:196        kwargs["environ_base"] = self._copy_environ(kwargs.get("environ_base", {}))197        builder = EnvironBuilder(self.application, *args, **kwargs)198 199        try:200            return builder.get_request()201        finally:202            builder.close()203 204    def open(205        self,206        *args: t.Any,207        buffered: bool = False,208        follow_redirects: bool = False,209        **kwargs: t.Any,210    ) -> TestResponse:211        if args and isinstance(212            args[0], (werkzeug.test.EnvironBuilder, dict, BaseRequest)213        ):214            if isinstance(args[0], werkzeug.test.EnvironBuilder):215                builder = copy(args[0])216                builder.environ_base = self._copy_environ(builder.environ_base or {})  # type: ignore[arg-type]217                request = builder.get_request()218            elif isinstance(args[0], dict):219                request = EnvironBuilder.from_environ(220                    args[0], app=self.application, environ_base=self._copy_environ({})221                ).get_request()222            else:223                # isinstance(args[0], BaseRequest)224                request = copy(args[0])225                request.environ = self._copy_environ(request.environ)226        else:227            # request is None228            request = self._request_from_builder_args(args, kwargs)229 230        # Pop any previously preserved contexts. This prevents contexts231        # from being preserved across redirects or multiple requests232        # within a single block.233        self._context_stack.close()234 235        response = super().open(236            request,237            buffered=buffered,238            follow_redirects=follow_redirects,239        )240        response.json_module = self.application.json  # type: ignore[assignment]241 242        # Re-push contexts that were preserved during the request.243        for cm in self._new_contexts:244            self._context_stack.enter_context(cm)245 246        self._new_contexts.clear()247        return response248 249    def __enter__(self) -> FlaskClient:250        if self.preserve_context:251            raise RuntimeError("Cannot nest client invocations")252        self.preserve_context = True253        return self254 255    def __exit__(256        self,257        exc_type: type | None,258        exc_value: BaseException | None,259        tb: TracebackType | None,260    ) -> None:261        self.preserve_context = False262        self._context_stack.close()263 264 265class FlaskCliRunner(CliRunner):266    """A :class:`~click.testing.CliRunner` for testing a Flask app's267    CLI commands. Typically created using268    :meth:`~flask.Flask.test_cli_runner`. See :ref:`testing-cli`.269    """270 271    def __init__(self, app: Flask, **kwargs: t.Any) -> None:272        self.app = app273        super().__init__(**kwargs)274 275    def invoke(  # type: ignore276        self, cli: t.Any = None, args: t.Any = None, **kwargs: t.Any277    ) -> Result:278        """Invokes a CLI command in an isolated environment. See279        :meth:`CliRunner.invoke <click.testing.CliRunner.invoke>` for280        full method documentation. See :ref:`testing-cli` for examples.281 282        If the ``obj`` argument is not given, passes an instance of283        :class:`~flask.cli.ScriptInfo` that knows how to load the Flask284        app being tested.285 286        :param cli: Command object to invoke. Default is the app's287            :attr:`~flask.app.Flask.cli` group.288        :param args: List of strings to invoke the command with.289 290        :return: a :class:`~click.testing.Result` object.291        """292        if cli is None:293            cli = self.app.cli294 295        if "obj" not in kwargs:296            kwargs["obj"] = ScriptInfo(create_app=lambda: self.app)297 298        return super().invoke(cli, args, **kwargs)299