testing.py 9.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283
  1. # -*- coding: utf-8 -*-
  2. """
  3. flask.testing
  4. ~~~~~~~~~~~~~
  5. Implements test support helpers. This module is lazily imported
  6. and usually not used in production environments.
  7. :copyright: 2010 Pallets
  8. :license: BSD-3-Clause
  9. """
  10. import warnings
  11. from contextlib import contextmanager
  12. import werkzeug.test
  13. from click.testing import CliRunner
  14. from werkzeug.test import Client
  15. from werkzeug.urls import url_parse
  16. from . import _request_ctx_stack
  17. from .cli import ScriptInfo
  18. from .json import dumps as json_dumps
  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,
  42. path="/",
  43. base_url=None,
  44. subdomain=None,
  45. url_scheme=None,
  46. *args,
  47. **kwargs
  48. ):
  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 = "{0}.{1}".format(subdomain, http_host)
  59. if url_scheme is None:
  60. url_scheme = app.config["PREFERRED_URL_SCHEME"]
  61. url = url_parse(path)
  62. base_url = "{scheme}://{netloc}/{path}".format(
  63. scheme=url.scheme or url_scheme,
  64. netloc=url.netloc or http_host,
  65. path=app_root.lstrip("/"),
  66. )
  67. path = url.path
  68. if url.query:
  69. sep = b"?" if isinstance(url.query, bytes) else "?"
  70. path += sep + url.query
  71. self.app = app
  72. super(EnvironBuilder, self).__init__(path, base_url, *args, **kwargs)
  73. def json_dumps(self, obj, **kwargs):
  74. """Serialize ``obj`` to a JSON-formatted string.
  75. The serialization will be configured according to the config associated
  76. with this EnvironBuilder's ``app``.
  77. """
  78. kwargs.setdefault("app", self.app)
  79. return json_dumps(obj, **kwargs)
  80. def make_test_environ_builder(*args, **kwargs):
  81. """Create a :class:`flask.testing.EnvironBuilder`.
  82. .. deprecated: 1.1
  83. Will be removed in 2.0. Construct
  84. ``flask.testing.EnvironBuilder`` directly instead.
  85. """
  86. warnings.warn(
  87. DeprecationWarning(
  88. '"make_test_environ_builder()" is deprecated and will be'
  89. ' removed in 2.0. Construct "flask.testing.EnvironBuilder"'
  90. " directly instead."
  91. )
  92. )
  93. return EnvironBuilder(*args, **kwargs)
  94. class FlaskClient(Client):
  95. """Works like a regular Werkzeug test client but has some knowledge about
  96. how Flask works to defer the cleanup of the request context stack to the
  97. end of a ``with`` body when used in a ``with`` statement. For general
  98. information about how to use this class refer to
  99. :class:`werkzeug.test.Client`.
  100. .. versionchanged:: 0.12
  101. `app.test_client()` includes preset default environment, which can be
  102. set after instantiation of the `app.test_client()` object in
  103. `client.environ_base`.
  104. Basic usage is outlined in the :ref:`testing` chapter.
  105. """
  106. preserve_context = False
  107. def __init__(self, *args, **kwargs):
  108. super(FlaskClient, self).__init__(*args, **kwargs)
  109. self.environ_base = {
  110. "REMOTE_ADDR": "127.0.0.1",
  111. "HTTP_USER_AGENT": "werkzeug/" + werkzeug.__version__,
  112. }
  113. @contextmanager
  114. def session_transaction(self, *args, **kwargs):
  115. """When used in combination with a ``with`` statement this opens a
  116. session transaction. This can be used to modify the session that
  117. the test client uses. Once the ``with`` block is left the session is
  118. stored back.
  119. ::
  120. with client.session_transaction() as session:
  121. session['value'] = 42
  122. Internally this is implemented by going through a temporary test
  123. request context and since session handling could depend on
  124. request variables this function accepts the same arguments as
  125. :meth:`~flask.Flask.test_request_context` which are directly
  126. passed through.
  127. """
  128. if self.cookie_jar is None:
  129. raise RuntimeError(
  130. "Session transactions only make sense with cookies enabled."
  131. )
  132. app = self.application
  133. environ_overrides = kwargs.setdefault("environ_overrides", {})
  134. self.cookie_jar.inject_wsgi(environ_overrides)
  135. outer_reqctx = _request_ctx_stack.top
  136. with app.test_request_context(*args, **kwargs) as c:
  137. session_interface = app.session_interface
  138. sess = session_interface.open_session(app, c.request)
  139. if sess is None:
  140. raise RuntimeError(
  141. "Session backend did not open a session. Check the configuration"
  142. )
  143. # Since we have to open a new request context for the session
  144. # handling we want to make sure that we hide out own context
  145. # from the caller. By pushing the original request context
  146. # (or None) on top of this and popping it we get exactly that
  147. # behavior. It's important to not use the push and pop
  148. # methods of the actual request context object since that would
  149. # mean that cleanup handlers are called
  150. _request_ctx_stack.push(outer_reqctx)
  151. try:
  152. yield sess
  153. finally:
  154. _request_ctx_stack.pop()
  155. resp = app.response_class()
  156. if not session_interface.is_null_session(sess):
  157. session_interface.save_session(app, sess, resp)
  158. headers = resp.get_wsgi_headers(c.request.environ)
  159. self.cookie_jar.extract_wsgi(c.request.environ, headers)
  160. def open(self, *args, **kwargs):
  161. as_tuple = kwargs.pop("as_tuple", False)
  162. buffered = kwargs.pop("buffered", False)
  163. follow_redirects = kwargs.pop("follow_redirects", False)
  164. if (
  165. not kwargs
  166. and len(args) == 1
  167. and isinstance(args[0], (werkzeug.test.EnvironBuilder, dict))
  168. ):
  169. environ = self.environ_base.copy()
  170. if isinstance(args[0], werkzeug.test.EnvironBuilder):
  171. environ.update(args[0].get_environ())
  172. else:
  173. environ.update(args[0])
  174. environ["flask._preserve_context"] = self.preserve_context
  175. else:
  176. kwargs.setdefault("environ_overrides", {})[
  177. "flask._preserve_context"
  178. ] = self.preserve_context
  179. kwargs.setdefault("environ_base", self.environ_base)
  180. builder = EnvironBuilder(self.application, *args, **kwargs)
  181. try:
  182. environ = builder.get_environ()
  183. finally:
  184. builder.close()
  185. return Client.open(
  186. self,
  187. environ,
  188. as_tuple=as_tuple,
  189. buffered=buffered,
  190. follow_redirects=follow_redirects,
  191. )
  192. def __enter__(self):
  193. if self.preserve_context:
  194. raise RuntimeError("Cannot nest client invocations")
  195. self.preserve_context = True
  196. return self
  197. def __exit__(self, exc_type, exc_value, tb):
  198. self.preserve_context = False
  199. # Normally the request context is preserved until the next
  200. # request in the same thread comes. When the client exits we
  201. # want to clean up earlier. Pop request contexts until the stack
  202. # is empty or a non-preserved one is found.
  203. while True:
  204. top = _request_ctx_stack.top
  205. if top is not None and top.preserved:
  206. top.pop()
  207. else:
  208. break
  209. class FlaskCliRunner(CliRunner):
  210. """A :class:`~click.testing.CliRunner` for testing a Flask app's
  211. CLI commands. Typically created using
  212. :meth:`~flask.Flask.test_cli_runner`. See :ref:`testing-cli`.
  213. """
  214. def __init__(self, app, **kwargs):
  215. self.app = app
  216. super(FlaskCliRunner, self).__init__(**kwargs)
  217. def invoke(self, cli=None, args=None, **kwargs):
  218. """Invokes a CLI command in an isolated environment. See
  219. :meth:`CliRunner.invoke <click.testing.CliRunner.invoke>` for
  220. full method documentation. See :ref:`testing-cli` for examples.
  221. If the ``obj`` argument is not given, passes an instance of
  222. :class:`~flask.cli.ScriptInfo` that knows how to load the Flask
  223. app being tested.
  224. :param cli: Command object to invoke. Default is the app's
  225. :attr:`~flask.app.Flask.cli` group.
  226. :param args: List of strings to invoke the command with.
  227. :return: a :class:`~click.testing.Result` object.
  228. """
  229. if cli is None:
  230. cli = self.app.cli
  231. if "obj" not in kwargs:
  232. kwargs["obj"] = ScriptInfo(create_app=lambda: self.app)
  233. return super(FlaskCliRunner, self).invoke(cli, args, **kwargs)