1import typing as t
2from contextlib import contextmanager
3from copy import copy
4from types import TracebackType
5
6import werkzeug.test
7from click.testing import CliRunner
8from werkzeug.test import Client
9from werkzeug.urls import url_parse
10from werkzeug.wrappers import Request as BaseRequest
11
12from . import _request_ctx_stack
13from .cli import ScriptInfo
14from .json import dumps as json_dumps
15from .sessions import SessionMixin
16
17if t.TYPE_CHECKING:
18    from .app import Flask
19    from .wrappers import Response
20
21
22class EnvironBuilder(werkzeug.test.EnvironBuilder):
23    """An :class:`~werkzeug.test.EnvironBuilder`, that takes defaults from the
24    application.
25
26    :param app: The Flask application to configure the environment from.
27    :param path: URL path being requested.
28    :param base_url: Base URL where the app is being served, which
29        ``path`` is relative to. If not given, built from
30        :data:`PREFERRED_URL_SCHEME`, ``subdomain``,
31        :data:`SERVER_NAME`, and :data:`APPLICATION_ROOT`.
32    :param subdomain: Subdomain name to append to :data:`SERVER_NAME`.
33    :param url_scheme: Scheme to use instead of
34        :data:`PREFERRED_URL_SCHEME`.
35    :param json: If given, this is serialized as JSON and passed as
36        ``data``. Also defaults ``content_type`` to
37        ``application/json``.
38    :param args: other positional arguments passed to
39        :class:`~werkzeug.test.EnvironBuilder`.
40    :param kwargs: other keyword arguments passed to
41        :class:`~werkzeug.test.EnvironBuilder`.
42    """
43
44    def __init__(
45        self,
46        app: "Flask",
47        path: str = "/",
48        base_url: t.Optional[str] = None,
49        subdomain: t.Optional[str] = None,
50        url_scheme: t.Optional[str] = None,
51        *args: t.Any,
52        **kwargs: t.Any,
53    ) -> None:
54        assert not (base_url or subdomain or url_scheme) or (
55            base_url is not None
56        ) != bool(
57            subdomain or url_scheme
58        ), 'Cannot pass "subdomain" or "url_scheme" with "base_url".'
59
60        if base_url is None:
61            http_host = app.config.get("SERVER_NAME") or "localhost"
62            app_root = app.config["APPLICATION_ROOT"]
63
64            if subdomain:
65                http_host = f"{subdomain}.{http_host}"
66
67            if url_scheme is None:
68                url_scheme = app.config["PREFERRED_URL_SCHEME"]
69
70            url = url_parse(path)
71            base_url = (
72                f"{url.scheme or url_scheme}://{url.netloc or http_host}"
73                f"/{app_root.lstrip('/')}"
74            )
75            path = url.path
76
77            if url.query:
78                sep = b"?" if isinstance(url.query, bytes) else "?"
79                path += sep + url.query
80
81        self.app = app
82        super().__init__(path, base_url, *args, **kwargs)
83
84    def json_dumps(self, obj: t.Any, **kwargs: t.Any) -> str:  # type: ignore
85        """Serialize ``obj`` to a JSON-formatted string.
86
87        The serialization will be configured according to the config associated
88        with this EnvironBuilder's ``app``.
89        """
90        kwargs.setdefault("app", self.app)
91        return json_dumps(obj, **kwargs)
92
93
94class FlaskClient(Client):
95    """Works like a regular Werkzeug test client but has some knowledge about
96    how Flask works to defer the cleanup of the request context stack to the
97    end of a ``with`` body when used in a ``with`` statement.  For general
98    information about how to use this class refer to
99    :class:`werkzeug.test.Client`.
100
101    .. versionchanged:: 0.12
102       `app.test_client()` includes preset default environment, which can be
103       set after instantiation of the `app.test_client()` object in
104       `client.environ_base`.
105
106    Basic usage is outlined in the :doc:`/testing` chapter.
107    """
108
109    application: "Flask"
110    preserve_context = False
111
112    def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
113        super().__init__(*args, **kwargs)
114        self.environ_base = {
115            "REMOTE_ADDR": "127.0.0.1",
116            "HTTP_USER_AGENT": f"werkzeug/{werkzeug.__version__}",
117        }
118
119    @contextmanager
120    def session_transaction(
121        self, *args: t.Any, **kwargs: t.Any
122    ) -> t.Generator[SessionMixin, None, None]:
123        """When used in combination with a ``with`` statement this opens a
124        session transaction.  This can be used to modify the session that
125        the test client uses.  Once the ``with`` block is left the session is
126        stored back.
127
128        ::
129
130            with client.session_transaction() as session:
131                session['value'] = 42
132
133        Internally this is implemented by going through a temporary test
134        request context and since session handling could depend on
135        request variables this function accepts the same arguments as
136        :meth:`~flask.Flask.test_request_context` which are directly
137        passed through.
138        """
139        if self.cookie_jar is None:
140            raise RuntimeError(
141                "Session transactions only make sense with cookies enabled."
142            )
143        app = self.application
144        environ_overrides = kwargs.setdefault("environ_overrides", {})
145        self.cookie_jar.inject_wsgi(environ_overrides)
146        outer_reqctx = _request_ctx_stack.top
147        with app.test_request_context(*args, **kwargs) as c:
148            session_interface = app.session_interface
149            sess = session_interface.open_session(app, c.request)
150            if sess is None:
151                raise RuntimeError(
152                    "Session backend did not open a session. Check the configuration"
153                )
154
155            # Since we have to open a new request context for the session
156            # handling we want to make sure that we hide out own context
157            # from the caller.  By pushing the original request context
158            # (or None) on top of this and popping it we get exactly that
159            # behavior.  It's important to not use the push and pop
160            # methods of the actual request context object since that would
161            # mean that cleanup handlers are called
162            _request_ctx_stack.push(outer_reqctx)
163            try:
164                yield sess
165            finally:
166                _request_ctx_stack.pop()
167
168            resp = app.response_class()
169            if not session_interface.is_null_session(sess):
170                session_interface.save_session(app, sess, resp)
171            headers = resp.get_wsgi_headers(c.request.environ)
172            self.cookie_jar.extract_wsgi(c.request.environ, headers)
173
174    def open(  # type: ignore
175        self,
176        *args: t.Any,
177        as_tuple: bool = False,
178        buffered: bool = False,
179        follow_redirects: bool = False,
180        **kwargs: t.Any,
181    ) -> "Response":
182        # Same logic as super.open, but apply environ_base and preserve_context.
183        request = None
184
185        def copy_environ(other):
186            return {
187                **self.environ_base,
188                **other,
189                "flask._preserve_context": self.preserve_context,
190            }
191
192        if not kwargs and len(args) == 1:
193            arg = args[0]
194
195            if isinstance(arg, werkzeug.test.EnvironBuilder):
196                builder = copy(arg)
197                builder.environ_base = copy_environ(builder.environ_base or {})
198                request = builder.get_request()
199            elif isinstance(arg, dict):
200                request = EnvironBuilder.from_environ(
201                    arg, app=self.application, environ_base=copy_environ({})
202                ).get_request()
203            elif isinstance(arg, BaseRequest):
204                request = copy(arg)
205                request.environ = copy_environ(request.environ)
206
207        if request is None:
208            kwargs["environ_base"] = copy_environ(kwargs.get("environ_base", {}))
209            builder = EnvironBuilder(self.application, *args, **kwargs)
210
211            try:
212                request = builder.get_request()
213            finally:
214                builder.close()
215
216        return super().open(  # type: ignore
217            request,
218            as_tuple=as_tuple,
219            buffered=buffered,
220            follow_redirects=follow_redirects,
221        )
222
223    def __enter__(self) -> "FlaskClient":
224        if self.preserve_context:
225            raise RuntimeError("Cannot nest client invocations")
226        self.preserve_context = True
227        return self
228
229    def __exit__(
230        self, exc_type: type, exc_value: BaseException, tb: TracebackType
231    ) -> None:
232        self.preserve_context = False
233
234        # Normally the request context is preserved until the next
235        # request in the same thread comes. When the client exits we
236        # want to clean up earlier. Pop request contexts until the stack
237        # is empty or a non-preserved one is found.
238        while True:
239            top = _request_ctx_stack.top
240
241            if top is not None and top.preserved:
242                top.pop()
243            else:
244                break
245
246
247class FlaskCliRunner(CliRunner):
248    """A :class:`~click.testing.CliRunner` for testing a Flask app's
249    CLI commands. Typically created using
250    :meth:`~flask.Flask.test_cli_runner`. See :ref:`testing-cli`.
251    """
252
253    def __init__(self, app: "Flask", **kwargs: t.Any) -> None:
254        self.app = app
255        super().__init__(**kwargs)
256
257    def invoke(  # type: ignore
258        self, cli: t.Any = None, args: t.Any = None, **kwargs: t.Any
259    ) -> t.Any:
260        """Invokes a CLI command in an isolated environment. See
261        :meth:`CliRunner.invoke <click.testing.CliRunner.invoke>` for
262        full method documentation. See :ref:`testing-cli` for examples.
263
264        If the ``obj`` argument is not given, passes an instance of
265        :class:`~flask.cli.ScriptInfo` that knows how to load the Flask
266        app being tested.
267
268        :param cli: Command object to invoke. Default is the app's
269            :attr:`~flask.app.Flask.cli` group.
270        :param args: List of strings to invoke the command with.
271
272        :return: a :class:`~click.testing.Result` object.
273        """
274        if cli is None:
275            cli = self.app.cli
276
277        if "obj" not in kwargs:
278            kwargs["obj"] = ScriptInfo(create_app=lambda: self.app)
279
280        return super().invoke(cli, args, **kwargs)
281