testing.py 10 KB

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