testing.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129
  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: (c) 2015 by Armin Ronacher.
  8. :license: BSD, see LICENSE for more details.
  9. """
  10. from contextlib import contextmanager
  11. from werkzeug.test import Client, EnvironBuilder
  12. from flask import _request_ctx_stack
  13. try:
  14. from werkzeug.urls import url_parse
  15. except ImportError:
  16. from urlparse import urlsplit as url_parse
  17. def make_test_environ_builder(app, path='/', base_url=None, *args, **kwargs):
  18. """Creates a new test builder with some application defaults thrown in."""
  19. http_host = app.config.get('SERVER_NAME')
  20. app_root = app.config.get('APPLICATION_ROOT')
  21. if base_url is None:
  22. url = url_parse(path)
  23. base_url = 'http://%s/' % (url.netloc or http_host or 'localhost')
  24. if app_root:
  25. base_url += app_root.lstrip('/')
  26. if url.netloc:
  27. path = url.path
  28. if url.query:
  29. path += '?' + url.query
  30. return EnvironBuilder(path, base_url, *args, **kwargs)
  31. class FlaskClient(Client):
  32. """Works like a regular Werkzeug test client but has some knowledge about
  33. how Flask works to defer the cleanup of the request context stack to the
  34. end of a ``with`` body when used in a ``with`` statement. For general
  35. information about how to use this class refer to
  36. :class:`werkzeug.test.Client`.
  37. Basic usage is outlined in the :ref:`testing` chapter.
  38. """
  39. preserve_context = False
  40. @contextmanager
  41. def session_transaction(self, *args, **kwargs):
  42. """When used in combination with a ``with`` statement this opens a
  43. session transaction. This can be used to modify the session that
  44. the test client uses. Once the ``with`` block is left the session is
  45. stored back.
  46. ::
  47. with client.session_transaction() as session:
  48. session['value'] = 42
  49. Internally this is implemented by going through a temporary test
  50. request context and since session handling could depend on
  51. request variables this function accepts the same arguments as
  52. :meth:`~flask.Flask.test_request_context` which are directly
  53. passed through.
  54. """
  55. if self.cookie_jar is None:
  56. raise RuntimeError('Session transactions only make sense '
  57. 'with cookies enabled.')
  58. app = self.application
  59. environ_overrides = kwargs.setdefault('environ_overrides', {})
  60. self.cookie_jar.inject_wsgi(environ_overrides)
  61. outer_reqctx = _request_ctx_stack.top
  62. with app.test_request_context(*args, **kwargs) as c:
  63. sess = app.open_session(c.request)
  64. if sess is None:
  65. raise RuntimeError('Session backend did not open a session. '
  66. 'Check the configuration')
  67. # Since we have to open a new request context for the session
  68. # handling we want to make sure that we hide out own context
  69. # from the caller. By pushing the original request context
  70. # (or None) on top of this and popping it we get exactly that
  71. # behavior. It's important to not use the push and pop
  72. # methods of the actual request context object since that would
  73. # mean that cleanup handlers are called
  74. _request_ctx_stack.push(outer_reqctx)
  75. try:
  76. yield sess
  77. finally:
  78. _request_ctx_stack.pop()
  79. resp = app.response_class()
  80. if not app.session_interface.is_null_session(sess):
  81. app.save_session(sess, resp)
  82. headers = resp.get_wsgi_headers(c.request.environ)
  83. self.cookie_jar.extract_wsgi(c.request.environ, headers)
  84. def open(self, *args, **kwargs):
  85. kwargs.setdefault('environ_overrides', {}) \
  86. ['flask._preserve_context'] = self.preserve_context
  87. as_tuple = kwargs.pop('as_tuple', False)
  88. buffered = kwargs.pop('buffered', False)
  89. follow_redirects = kwargs.pop('follow_redirects', False)
  90. builder = make_test_environ_builder(self.application, *args, **kwargs)
  91. return Client.open(self, builder,
  92. as_tuple=as_tuple,
  93. buffered=buffered,
  94. follow_redirects=follow_redirects)
  95. def __enter__(self):
  96. if self.preserve_context:
  97. raise RuntimeError('Cannot nest client invocations')
  98. self.preserve_context = True
  99. return self
  100. def __exit__(self, exc_type, exc_value, tb):
  101. self.preserve_context = False
  102. # on exit we want to clean up earlier. Normally the request context
  103. # stays preserved until the next request in the same thread comes
  104. # in. See RequestGlobals.push() for the general behavior.
  105. top = _request_ctx_stack.top
  106. if top is not None and top.preserved:
  107. top.pop()