testing.py 4.9 KB

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