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