1  from __future__ import annotations
   2  
   3  import importlib.metadata
   4  import typing as t
   5  from contextlib import contextmanager
   6  from contextlib import ExitStack
   7  from copy import copy
   8  from types import TracebackType
   9  from urllib.parse import urlsplit
  10  
  11  import werkzeug.test
  12  from click.testing import CliRunner
  13  from werkzeug.test import Client
  14  from werkzeug.wrappers import Request as BaseRequest
  15  
  16  from .cli import ScriptInfo
  17  from .sessions import SessionMixin
  18  
  19  if t.TYPE_CHECKING:  # pragma: no cover
  20      from _typeshed.wsgi import WSGIEnvironment
  21      from werkzeug.test import TestResponse
  22  
  23      from .app import Flask
  24  
  25  
  26  class EnvironBuilder(werkzeug.test.EnvironBuilder):
  27      """An :class:`~werkzeug.test.EnvironBuilder`, that takes defaults from the
  28      application.
  29  
  30      :param app: The Flask application to configure the environment from.
  31      :param path: URL path being requested.
  32      :param base_url: Base URL where the app is being served, which
  33          ``path`` is relative to. If not given, built from
  34          :data:`PREFERRED_URL_SCHEME`, ``subdomain``,
  35          :data:`SERVER_NAME`, and :data:`APPLICATION_ROOT`.
  36      :param subdomain: Subdomain name to append to :data:`SERVER_NAME`.
  37      :param url_scheme: Scheme to use instead of
  38          :data:`PREFERRED_URL_SCHEME`.
  39      :param json: If given, this is serialized as JSON and passed as
  40          ``data``. Also defaults ``content_type`` to
  41          ``application/json``.
  42      :param args: other positional arguments passed to
  43          :class:`~werkzeug.test.EnvironBuilder`.
  44      :param kwargs: other keyword arguments passed to
  45          :class:`~werkzeug.test.EnvironBuilder`.
  46      """
  47  
  48      def __init__(
  49          self,
  50          app: Flask,
  51          path: str = "/",
  52          base_url: str | None = None,
  53          subdomain: str | None = None,
  54          url_scheme: str | None = None,
  55          *args: t.Any,
  56          **kwargs: t.Any,
  57      ) -> None:
  58          assert not (base_url or subdomain or url_scheme) or (
  59              base_url is not None
  60          ) != bool(
  61              subdomain or url_scheme
  62          ), 'Cannot pass "subdomain" or "url_scheme" with "base_url".'
  63  
  64          if base_url is None:
  65              http_host = app.config.get("SERVER_NAME") or "localhost"
  66              app_root = app.config["APPLICATION_ROOT"]
  67  
  68              if subdomain:
  69                  http_host = f"{subdomain}.{http_host}"
  70  
  71              if url_scheme is None:
  72                  url_scheme = app.config["PREFERRED_URL_SCHEME"]
  73  
  74              url = urlsplit(path)
  75              base_url = (
  76                  f"{url.scheme or url_scheme}://{url.netloc or http_host}"
  77                  f"/{app_root.lstrip('/')}"
  78              )
  79              path = url.path
  80  
  81              if url.query:
  82                  path = f"{path}?{url.query}"
  83  
  84          self.app = app
  85          super().__init__(path, base_url, *args, **kwargs)
  86  
  87      def json_dumps(self, obj: t.Any, **kwargs: t.Any) -> str:  # type: ignore
  88          """Serialize ``obj`` to a JSON-formatted string.
  89  
  90          The serialization will be configured according to the config associated
  91          with this EnvironBuilder's ``app``.
  92          """
  93          return self.app.json.dumps(obj, **kwargs)
  94  
  95  
  96  _werkzeug_version = ""
  97  
  98  
  99  def _get_werkzeug_version() -> str:
 100      global _werkzeug_version
 101  
 102      if not _werkzeug_version:
 103          _werkzeug_version = importlib.metadata.version("werkzeug")
 104  
 105      return _werkzeug_version
 106  
 107  
 108  class FlaskClient(Client):
 109      """Works like a regular Werkzeug test client but has knowledge about
 110      Flask's contexts to defer the cleanup of the request context until
 111      the end of a ``with`` block. For general information about how to
 112      use this class refer to :class:`werkzeug.test.Client`.
 113  
 114      .. versionchanged:: 0.12
 115         `app.test_client()` includes preset default environment, which can be
 116         set after instantiation of the `app.test_client()` object in
 117         `client.environ_base`.
 118  
 119      Basic usage is outlined in the :doc:`/testing` chapter.
 120      """
 121  
 122      application: Flask
 123  
 124      def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
 125          super().__init__(*args, **kwargs)
 126          self.preserve_context = False
 127          self._new_contexts: list[t.ContextManager[t.Any]] = []
 128          self._context_stack = ExitStack()
 129          self.environ_base = {
 130              "REMOTE_ADDR": "127.0.0.1",
 131              "HTTP_USER_AGENT": f"Werkzeug/{_get_werkzeug_version()}",
 132          }
 133  
 134      @contextmanager
 135      def session_transaction(
 136          self, *args: t.Any, **kwargs: t.Any
 137      ) -> t.Iterator[SessionMixin]:
 138          """When used in combination with a ``with`` statement this opens a
 139          session transaction.  This can be used to modify the session that
 140          the test client uses.  Once the ``with`` block is left the session is
 141          stored back.
 142  
 143          ::
 144  
 145              with client.session_transaction() as session:
 146                  session['value'] = 42
 147  
 148          Internally this is implemented by going through a temporary test
 149          request context and since session handling could depend on
 150          request variables this function accepts the same arguments as
 151          :meth:`~flask.Flask.test_request_context` which are directly
 152          passed through.
 153          """
 154          if self._cookies is None:
 155              raise TypeError(
 156                  "Cookies are disabled. Create a client with 'use_cookies=True'."
 157              )
 158  
 159          app = self.application
 160          ctx = app.test_request_context(*args, **kwargs)
 161          self._add_cookies_to_wsgi(ctx.request.environ)
 162  
 163          with ctx:
 164              sess = app.session_interface.open_session(app, ctx.request)
 165  
 166          if sess is None:
 167              raise RuntimeError("Session backend did not open a session.")
 168  
 169          yield sess
 170          resp = app.response_class()
 171  
 172          if app.session_interface.is_null_session(sess):
 173              return
 174  
 175          with ctx:
 176              app.session_interface.save_session(app, sess, resp)
 177  
 178          self._update_cookies_from_response(
 179              ctx.request.host.partition(":")[0],
 180              ctx.request.path,
 181              resp.headers.getlist("Set-Cookie"),
 182          )
 183  
 184      def _copy_environ(self, other: WSGIEnvironment) -> WSGIEnvironment:
 185          out = {**self.environ_base, **other}
 186  
 187          if self.preserve_context:
 188              out["werkzeug.debug.preserve_context"] = self._new_contexts.append
 189  
 190          return out
 191  
 192      def _request_from_builder_args(
 193          self, args: tuple[t.Any, ...], kwargs: dict[str, t.Any]
 194      ) -> BaseRequest:
 195          kwargs["environ_base"] = self._copy_environ(kwargs.get("environ_base", {}))
 196          builder = EnvironBuilder(self.application, *args, **kwargs)
 197  
 198          try:
 199              return builder.get_request()
 200          finally:
 201              builder.close()
 202  
 203      def open(
 204          self,
 205          *args: t.Any,
 206          buffered: bool = False,
 207          follow_redirects: bool = False,
 208          **kwargs: t.Any,
 209      ) -> TestResponse:
 210          if args and isinstance(
 211              args[0], (werkzeug.test.EnvironBuilder, dict, BaseRequest)
 212          ):
 213              if isinstance(args[0], werkzeug.test.EnvironBuilder):
 214                  builder = copy(args[0])
 215                  builder.environ_base = self._copy_environ(builder.environ_base or {})  # type: ignore[arg-type]
 216                  request = builder.get_request()
 217              elif isinstance(args[0], dict):
 218                  request = EnvironBuilder.from_environ(
 219                      args[0], app=self.application, environ_base=self._copy_environ({})
 220                  ).get_request()
 221              else:
 222                  # isinstance(args[0], BaseRequest)
 223                  request = copy(args[0])
 224                  request.environ = self._copy_environ(request.environ)
 225          else:
 226              # request is None
 227              request = self._request_from_builder_args(args, kwargs)
 228  
 229          # Pop any previously preserved contexts. This prevents contexts
 230          # from being preserved across redirects or multiple requests
 231          # within a single block.
 232          self._context_stack.close()
 233  
 234          response = super().open(
 235              request,
 236              buffered=buffered,
 237              follow_redirects=follow_redirects,
 238          )
 239          response.json_module = self.application.json  # type: ignore[assignment]
 240  
 241          # Re-push contexts that were preserved during the request.
 242          while self._new_contexts:
 243              cm = self._new_contexts.pop()
 244              self._context_stack.enter_context(cm)
 245  
 246          return response
 247  
 248      def __enter__(self) -> FlaskClient:
 249          if self.preserve_context:
 250              raise RuntimeError("Cannot nest client invocations")
 251          self.preserve_context = True
 252          return self
 253  
 254      def __exit__(
 255          self,
 256          exc_type: type | None,
 257          exc_value: BaseException | None,
 258          tb: TracebackType | None,
 259      ) -> None:
 260          self.preserve_context = False
 261          self._context_stack.close()
 262  
 263  
 264  class FlaskCliRunner(CliRunner):
 265      """A :class:`~click.testing.CliRunner` for testing a Flask app's
 266      CLI commands. Typically created using
 267      :meth:`~flask.Flask.test_cli_runner`. See :ref:`testing-cli`.
 268      """
 269  
 270      def __init__(self, app: Flask, **kwargs: t.Any) -> None:
 271          self.app = app
 272          super().__init__(**kwargs)
 273  
 274      def invoke(  # type: ignore
 275          self, cli: t.Any = None, args: t.Any = None, **kwargs: t.Any
 276      ) -> t.Any:
 277          """Invokes a CLI command in an isolated environment. See
 278          :meth:`CliRunner.invoke <click.testing.CliRunner.invoke>` for
 279          full method documentation. See :ref:`testing-cli` for examples.
 280  
 281          If the ``obj`` argument is not given, passes an instance of
 282          :class:`~flask.cli.ScriptInfo` that knows how to load the Flask
 283          app being tested.
 284  
 285          :param cli: Command object to invoke. Default is the app's
 286              :attr:`~flask.app.Flask.cli` group.
 287          :param args: List of strings to invoke the command with.
 288  
 289          :return: a :class:`~click.testing.Result` object.
 290          """
 291          if cli is None:
 292              cli = self.app.cli
 293  
 294          if "obj" not in kwargs:
 295              kwargs["obj"] = ScriptInfo(create_app=lambda: self.app)
 296  
 297          return super().invoke(cli, args, **kwargs)
 298