testing.py 9.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250
  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 by the Pallets team.
  8. :license: BSD, see LICENSE for more details.
  9. """
  10. import werkzeug
  11. from contextlib import contextmanager
  12. from click.testing import CliRunner
  13. from flask.cli import ScriptInfo
  14. from werkzeug.test import Client, EnvironBuilder
  15. from flask import _request_ctx_stack
  16. from flask.json import dumps as json_dumps
  17. from werkzeug.urls import url_parse
  18. def make_test_environ_builder(
  19. app, path='/', base_url=None, subdomain=None, url_scheme=None,
  20. *args, **kwargs
  21. ):
  22. """Create a :class:`~werkzeug.test.EnvironBuilder`, taking some
  23. defaults from the application.
  24. :param app: The Flask application to configure the environment from.
  25. :param path: URL path being requested.
  26. :param base_url: Base URL where the app is being served, which
  27. ``path`` is relative to. If not given, built from
  28. :data:`PREFERRED_URL_SCHEME`, ``subdomain``,
  29. :data:`SERVER_NAME`, and :data:`APPLICATION_ROOT`.
  30. :param subdomain: Subdomain name to append to :data:`SERVER_NAME`.
  31. :param url_scheme: Scheme to use instead of
  32. :data:`PREFERRED_URL_SCHEME`.
  33. :param json: If given, this is serialized as JSON and passed as
  34. ``data``. Also defaults ``content_type`` to
  35. ``application/json``.
  36. :param args: other positional arguments passed to
  37. :class:`~werkzeug.test.EnvironBuilder`.
  38. :param kwargs: other keyword arguments passed to
  39. :class:`~werkzeug.test.EnvironBuilder`.
  40. """
  41. assert (
  42. not (base_url or subdomain or url_scheme)
  43. or (base_url is not None) != bool(subdomain or url_scheme)
  44. ), 'Cannot pass "subdomain" or "url_scheme" with "base_url".'
  45. if base_url is None:
  46. http_host = app.config.get('SERVER_NAME') or 'localhost'
  47. app_root = app.config['APPLICATION_ROOT']
  48. if subdomain:
  49. http_host = '{0}.{1}'.format(subdomain, http_host)
  50. if url_scheme is None:
  51. url_scheme = app.config['PREFERRED_URL_SCHEME']
  52. url = url_parse(path)
  53. base_url = '{scheme}://{netloc}/{path}'.format(
  54. scheme=url.scheme or url_scheme,
  55. netloc=url.netloc or http_host,
  56. path=app_root.lstrip('/')
  57. )
  58. path = url.path
  59. if url.query:
  60. sep = b'?' if isinstance(url.query, bytes) else '?'
  61. path += sep + url.query
  62. if 'json' in kwargs:
  63. assert 'data' not in kwargs, (
  64. "Client cannot provide both 'json' and 'data'."
  65. )
  66. # push a context so flask.json can use app's json attributes
  67. with app.app_context():
  68. kwargs['data'] = json_dumps(kwargs.pop('json'))
  69. if 'content_type' not in kwargs:
  70. kwargs['content_type'] = 'application/json'
  71. return EnvironBuilder(path, base_url, *args, **kwargs)
  72. class FlaskClient(Client):
  73. """Works like a regular Werkzeug test client but has some knowledge about
  74. how Flask works to defer the cleanup of the request context stack to the
  75. end of a ``with`` body when used in a ``with`` statement. For general
  76. information about how to use this class refer to
  77. :class:`werkzeug.test.Client`.
  78. .. versionchanged:: 0.12
  79. `app.test_client()` includes preset default environment, which can be
  80. set after instantiation of the `app.test_client()` object in
  81. `client.environ_base`.
  82. Basic usage is outlined in the :ref:`testing` chapter.
  83. """
  84. preserve_context = False
  85. def __init__(self, *args, **kwargs):
  86. super(FlaskClient, self).__init__(*args, **kwargs)
  87. self.environ_base = {
  88. "REMOTE_ADDR": "127.0.0.1",
  89. "HTTP_USER_AGENT": "werkzeug/" + werkzeug.__version__
  90. }
  91. @contextmanager
  92. def session_transaction(self, *args, **kwargs):
  93. """When used in combination with a ``with`` statement this opens a
  94. session transaction. This can be used to modify the session that
  95. the test client uses. Once the ``with`` block is left the session is
  96. stored back.
  97. ::
  98. with client.session_transaction() as session:
  99. session['value'] = 42
  100. Internally this is implemented by going through a temporary test
  101. request context and since session handling could depend on
  102. request variables this function accepts the same arguments as
  103. :meth:`~flask.Flask.test_request_context` which are directly
  104. passed through.
  105. """
  106. if self.cookie_jar is None:
  107. raise RuntimeError('Session transactions only make sense '
  108. 'with cookies enabled.')
  109. app = self.application
  110. environ_overrides = kwargs.setdefault('environ_overrides', {})
  111. self.cookie_jar.inject_wsgi(environ_overrides)
  112. outer_reqctx = _request_ctx_stack.top
  113. with app.test_request_context(*args, **kwargs) as c:
  114. session_interface = app.session_interface
  115. sess = session_interface.open_session(app, c.request)
  116. if sess is None:
  117. raise RuntimeError('Session backend did not open a session. '
  118. 'Check the configuration')
  119. # Since we have to open a new request context for the session
  120. # handling we want to make sure that we hide out own context
  121. # from the caller. By pushing the original request context
  122. # (or None) on top of this and popping it we get exactly that
  123. # behavior. It's important to not use the push and pop
  124. # methods of the actual request context object since that would
  125. # mean that cleanup handlers are called
  126. _request_ctx_stack.push(outer_reqctx)
  127. try:
  128. yield sess
  129. finally:
  130. _request_ctx_stack.pop()
  131. resp = app.response_class()
  132. if not session_interface.is_null_session(sess):
  133. session_interface.save_session(app, sess, resp)
  134. headers = resp.get_wsgi_headers(c.request.environ)
  135. self.cookie_jar.extract_wsgi(c.request.environ, headers)
  136. def open(self, *args, **kwargs):
  137. as_tuple = kwargs.pop('as_tuple', False)
  138. buffered = kwargs.pop('buffered', False)
  139. follow_redirects = kwargs.pop('follow_redirects', False)
  140. if (
  141. not kwargs and len(args) == 1
  142. and isinstance(args[0], (EnvironBuilder, dict))
  143. ):
  144. environ = self.environ_base.copy()
  145. if isinstance(args[0], EnvironBuilder):
  146. environ.update(args[0].get_environ())
  147. else:
  148. environ.update(args[0])
  149. environ['flask._preserve_context'] = self.preserve_context
  150. else:
  151. kwargs.setdefault('environ_overrides', {}) \
  152. ['flask._preserve_context'] = self.preserve_context
  153. kwargs.setdefault('environ_base', self.environ_base)
  154. builder = make_test_environ_builder(
  155. self.application, *args, **kwargs
  156. )
  157. try:
  158. environ = builder.get_environ()
  159. finally:
  160. builder.close()
  161. return Client.open(
  162. self, environ,
  163. as_tuple=as_tuple,
  164. buffered=buffered,
  165. follow_redirects=follow_redirects
  166. )
  167. def __enter__(self):
  168. if self.preserve_context:
  169. raise RuntimeError('Cannot nest client invocations')
  170. self.preserve_context = True
  171. return self
  172. def __exit__(self, exc_type, exc_value, tb):
  173. self.preserve_context = False
  174. # on exit we want to clean up earlier. Normally the request context
  175. # stays preserved until the next request in the same thread comes
  176. # in. See RequestGlobals.push() for the general behavior.
  177. top = _request_ctx_stack.top
  178. if top is not None and top.preserved:
  179. top.pop()
  180. class FlaskCliRunner(CliRunner):
  181. """A :class:`~click.testing.CliRunner` for testing a Flask app's
  182. CLI commands. Typically created using
  183. :meth:`~flask.Flask.test_cli_runner`. See :ref:`testing-cli`.
  184. """
  185. def __init__(self, app, **kwargs):
  186. self.app = app
  187. super(FlaskCliRunner, self).__init__(**kwargs)
  188. def invoke(self, cli=None, args=None, **kwargs):
  189. """Invokes a CLI command in an isolated environment. See
  190. :meth:`CliRunner.invoke <click.testing.CliRunner.invoke>` for
  191. full method documentation. See :ref:`testing-cli` for examples.
  192. If the ``obj`` argument is not given, passes an instance of
  193. :class:`~flask.cli.ScriptInfo` that knows how to load the Flask
  194. app being tested.
  195. :param cli: Command object to invoke. Default is the app's
  196. :attr:`~flask.app.Flask.cli` group.
  197. :param args: List of strings to invoke the command with.
  198. :return: a :class:`~click.testing.Result` object.
  199. """
  200. if cli is None:
  201. cli = self.app.cli
  202. if 'obj' not in kwargs:
  203. kwargs['obj'] = ScriptInfo(create_app=lambda: self.app)
  204. return super(FlaskCliRunner, self).invoke(cli, args, **kwargs)