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