testing.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283
  1. import typing as t
  2. from contextlib import contextmanager
  3. from copy import copy
  4. from types import TracebackType
  5. import werkzeug.test
  6. from click.testing import CliRunner
  7. from werkzeug.test import Client
  8. from werkzeug.urls import url_parse
  9. from werkzeug.wrappers import Request as BaseRequest
  10. from .cli import ScriptInfo
  11. from .globals import _request_ctx_stack
  12. from .json import dumps as json_dumps
  13. from .sessions import SessionMixin
  14. if t.TYPE_CHECKING:
  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. kwargs.setdefault("app", self.app)
  76. return json_dumps(obj, **kwargs)
  77. class FlaskClient(Client):
  78. """Works like a regular Werkzeug test client but has some knowledge about
  79. how Flask works to defer the cleanup of the request context stack to the
  80. end of a ``with`` body when used in a ``with`` statement. For general
  81. information about how to use this class refer to
  82. :class:`werkzeug.test.Client`.
  83. .. versionchanged:: 0.12
  84. `app.test_client()` includes preset default environment, which can be
  85. set after instantiation of the `app.test_client()` object in
  86. `client.environ_base`.
  87. Basic usage is outlined in the :doc:`/testing` chapter.
  88. """
  89. application: "Flask"
  90. preserve_context = False
  91. def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
  92. super().__init__(*args, **kwargs)
  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 = _request_ctx_stack.top
  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. _request_ctx_stack.push(outer_reqctx)
  137. try:
  138. yield sess
  139. finally:
  140. _request_ctx_stack.pop()
  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. return {
  148. **self.environ_base,
  149. **other,
  150. "flask._preserve_context": self.preserve_context,
  151. }
  152. def _request_from_builder_args(self, args, kwargs):
  153. kwargs["environ_base"] = self._copy_environ(kwargs.get("environ_base", {}))
  154. builder = EnvironBuilder(self.application, *args, **kwargs)
  155. try:
  156. return builder.get_request()
  157. finally:
  158. builder.close()
  159. def open(
  160. self,
  161. *args: t.Any,
  162. buffered: bool = False,
  163. follow_redirects: bool = False,
  164. **kwargs: t.Any,
  165. ) -> "TestResponse":
  166. if args and isinstance(
  167. args[0], (werkzeug.test.EnvironBuilder, dict, BaseRequest)
  168. ):
  169. if isinstance(args[0], werkzeug.test.EnvironBuilder):
  170. builder = copy(args[0])
  171. builder.environ_base = self._copy_environ(builder.environ_base or {})
  172. request = builder.get_request()
  173. elif isinstance(args[0], dict):
  174. request = EnvironBuilder.from_environ(
  175. args[0], app=self.application, environ_base=self._copy_environ({})
  176. ).get_request()
  177. else:
  178. # isinstance(args[0], BaseRequest)
  179. request = copy(args[0])
  180. request.environ = self._copy_environ(request.environ)
  181. else:
  182. # request is None
  183. request = self._request_from_builder_args(args, kwargs)
  184. return super().open(
  185. request,
  186. buffered=buffered,
  187. follow_redirects=follow_redirects,
  188. )
  189. def __enter__(self) -> "FlaskClient":
  190. if self.preserve_context:
  191. raise RuntimeError("Cannot nest client invocations")
  192. self.preserve_context = True
  193. return self
  194. def __exit__(
  195. self,
  196. exc_type: t.Optional[type],
  197. exc_value: t.Optional[BaseException],
  198. tb: t.Optional[TracebackType],
  199. ) -> None:
  200. self.preserve_context = False
  201. # Normally the request context is preserved until the next
  202. # request in the same thread comes. When the client exits we
  203. # want to clean up earlier. Pop request contexts until the stack
  204. # is empty or a non-preserved one is found.
  205. while True:
  206. top = _request_ctx_stack.top
  207. if top is not None and top.preserved:
  208. top.pop()
  209. else:
  210. break
  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)