testing.py 9.9 KB

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