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