import typing as t from contextlib import contextmanager from contextlib import ExitStack from copy import copy from types import TracebackType import werkzeug.test from click.testing import CliRunner from werkzeug.test import Client from werkzeug.urls import url_parse from werkzeug.wrappers import Request as BaseRequest from .cli import ScriptInfo from .globals import _cv_request from .sessions import SessionMixin if t.TYPE_CHECKING: # pragma: no cover from werkzeug.test import TestResponse from .app import Flask class EnvironBuilder(werkzeug.test.EnvironBuilder): """An :class:`~werkzeug.test.EnvironBuilder`, that takes defaults from the application. :param app: The Flask application to configure the environment from. :param path: URL path being requested. :param base_url: Base URL where the app is being served, which ``path`` is relative to. If not given, built from :data:`PREFERRED_URL_SCHEME`, ``subdomain``, :data:`SERVER_NAME`, and :data:`APPLICATION_ROOT`. :param subdomain: Subdomain name to append to :data:`SERVER_NAME`. :param url_scheme: Scheme to use instead of :data:`PREFERRED_URL_SCHEME`. :param json: If given, this is serialized as JSON and passed as ``data``. Also defaults ``content_type`` to ``application/json``. :param args: other positional arguments passed to :class:`~werkzeug.test.EnvironBuilder`. :param kwargs: other keyword arguments passed to :class:`~werkzeug.test.EnvironBuilder`. """ def __init__( self, app: "Flask", path: str = "/", base_url: t.Optional[str] = None, subdomain: t.Optional[str] = None, url_scheme: t.Optional[str] = None, *args: t.Any, **kwargs: t.Any, ) -> None: assert not (base_url or subdomain or url_scheme) or ( base_url is not None ) != bool( subdomain or url_scheme ), 'Cannot pass "subdomain" or "url_scheme" with "base_url".' if base_url is None: http_host = app.config.get("SERVER_NAME") or "localhost" app_root = app.config["APPLICATION_ROOT"] if subdomain: http_host = f"{subdomain}.{http_host}" if url_scheme is None: url_scheme = app.config["PREFERRED_URL_SCHEME"] url = url_parse(path) base_url = ( f"{url.scheme or url_scheme}://{url.netloc or http_host}" f"/{app_root.lstrip('/')}" ) path = url.path if url.query: sep = b"?" if isinstance(url.query, bytes) else "?" path += sep + url.query self.app = app super().__init__(path, base_url, *args, **kwargs) def json_dumps(self, obj: t.Any, **kwargs: t.Any) -> str: # type: ignore """Serialize ``obj`` to a JSON-formatted string. The serialization will be configured according to the config associated with this EnvironBuilder's ``app``. """ return self.app.json.dumps(obj, **kwargs) class FlaskClient(Client): """Works like a regular Werkzeug test client but has knowledge about Flask's contexts to defer the cleanup of the request context until the end of a ``with`` block. For general information about how to use this class refer to :class:`werkzeug.test.Client`. .. versionchanged:: 0.12 `app.test_client()` includes preset default environment, which can be set after instantiation of the `app.test_client()` object in `client.environ_base`. Basic usage is outlined in the :doc:`/testing` chapter. """ application: "Flask" def __init__(self, *args: t.Any, **kwargs: t.Any) -> None: super().__init__(*args, **kwargs) self.preserve_context = False self._new_contexts: t.List[t.ContextManager[t.Any]] = [] self._context_stack = ExitStack() self.environ_base = { "REMOTE_ADDR": "127.0.0.1", "HTTP_USER_AGENT": f"werkzeug/{werkzeug.__version__}", } @contextmanager def session_transaction( self, *args: t.Any, **kwargs: t.Any ) -> t.Generator[SessionMixin, None, None]: """When used in combination with a ``with`` statement this opens a session transaction. This can be used to modify the session that the test client uses. Once the ``with`` block is left the session is stored back. :: with client.session_transaction() as session: session['value'] = 42 Internally this is implemented by going through a temporary test request context and since session handling could depend on request variables this function accepts the same arguments as :meth:`~flask.Flask.test_request_context` which are directly passed through. """ if self.cookie_jar is None: raise RuntimeError( "Session transactions only make sense with cookies enabled." ) app = self.application environ_overrides = kwargs.setdefault("environ_overrides", {}) self.cookie_jar.inject_wsgi(environ_overrides) outer_reqctx = _cv_request.get(None) with app.test_request_context(*args, **kwargs) as c: session_interface = app.session_interface sess = session_interface.open_session(app, c.request) if sess is None: raise RuntimeError( "Session backend did not open a session. Check the configuration" ) # Since we have to open a new request context for the session # handling we want to make sure that we hide out own context # from the caller. By pushing the original request context # (or None) on top of this and popping it we get exactly that # behavior. It's important to not use the push and pop # methods of the actual request context object since that would # mean that cleanup handlers are called token = _cv_request.set(outer_reqctx) # type: ignore[arg-type] try: yield sess finally: _cv_request.reset(token) resp = app.response_class() if not session_interface.is_null_session(sess): session_interface.save_session(app, sess, resp) headers = resp.get_wsgi_headers(c.request.environ) self.cookie_jar.extract_wsgi(c.request.environ, headers) def _copy_environ(self, other): out = {**self.environ_base, **other} if self.preserve_context: out["werkzeug.debug.preserve_context"] = self._new_contexts.append return out def _request_from_builder_args(self, args, kwargs): kwargs["environ_base"] = self._copy_environ(kwargs.get("environ_base", {})) builder = EnvironBuilder(self.application, *args, **kwargs) try: return builder.get_request() finally: builder.close() def open( self, *args: t.Any, buffered: bool = False, follow_redirects: bool = False, **kwargs: t.Any, ) -> "TestResponse": if args and isinstance( args[0], (werkzeug.test.EnvironBuilder, dict, BaseRequest) ): if isinstance(args[0], werkzeug.test.EnvironBuilder): builder = copy(args[0]) builder.environ_base = self._copy_environ(builder.environ_base or {}) request = builder.get_request() elif isinstance(args[0], dict): request = EnvironBuilder.from_environ( args[0], app=self.application, environ_base=self._copy_environ({}) ).get_request() else: # isinstance(args[0], BaseRequest) request = copy(args[0]) request.environ = self._copy_environ(request.environ) else: # request is None request = self._request_from_builder_args(args, kwargs) # Pop any previously preserved contexts. This prevents contexts # from being preserved across redirects or multiple requests # within a single block. self._context_stack.close() response = super().open( request, buffered=buffered, follow_redirects=follow_redirects, ) response.json_module = self.application.json # type: ignore[misc] # Re-push contexts that were preserved during the request. while self._new_contexts: cm = self._new_contexts.pop() self._context_stack.enter_context(cm) return response def __enter__(self) -> "FlaskClient": if self.preserve_context: raise RuntimeError("Cannot nest client invocations") self.preserve_context = True return self def __exit__( self, exc_type: t.Optional[type], exc_value: t.Optional[BaseException], tb: t.Optional[TracebackType], ) -> None: self.preserve_context = False self._context_stack.close() class FlaskCliRunner(CliRunner): """A :class:`~click.testing.CliRunner` for testing a Flask app's CLI commands. Typically created using :meth:`~flask.Flask.test_cli_runner`. See :ref:`testing-cli`. """ def __init__(self, app: "Flask", **kwargs: t.Any) -> None: self.app = app super().__init__(**kwargs) def invoke( # type: ignore self, cli: t.Any = None, args: t.Any = None, **kwargs: t.Any ) -> t.Any: """Invokes a CLI command in an isolated environment. See :meth:`CliRunner.invoke ` for full method documentation. See :ref:`testing-cli` for examples. If the ``obj`` argument is not given, passes an instance of :class:`~flask.cli.ScriptInfo` that knows how to load the Flask app being tested. :param cli: Command object to invoke. Default is the app's :attr:`~flask.app.Flask.cli` group. :param args: List of strings to invoke the command with. :return: a :class:`~click.testing.Result` object. """ if cli is None: cli = self.app.cli # type: ignore if "obj" not in kwargs: kwargs["obj"] = ScriptInfo(create_app=lambda: self.app) return super().invoke(cli, args, **kwargs)