Implemented flask.testing.TestClient.session_transaction for quick session modifications in test environments.

CHANGES

@@ -33,6 +33,8 @@ Relase date to be decided, codename to be chosen.
   the perfect place to put configuration files etc.  For more information
   see :ref:`instance-folders`.
 - Added the ``APPLICATION_ROOT`` configuration variable.
+- Implemented :meth:`~flask.testing.TestClient.session_transaction` to
+  easily modify sessions from the test environment.
 
 Version 0.7.3
 -------------

docs/api.rst

@@ -218,6 +218,15 @@ implementation that Flask is using.
    :members:
 
 
+Test Client
+-----------
+
+.. currentmodule:: flask.testing
+
+.. autoclass:: TestClient
+   :members:
+
+
 Application Globals
 -------------------
 

flask/app.py

@@ -706,6 +706,8 @@ class Flask(_PackageBoundObject):
                 rv = c.get('/?vodka=42')
                 assert request.args['vodka'] == '42'
 
+        See :class:`~flask.testing.TestClient` for more information.
+
         .. versionchanged:: 0.4
            added support for `with` block usage for the client.
 

flask/testing.py

@@ -10,19 +10,69 @@
     :license: BSD, see LICENSE for more details.
 """
 
+from contextlib import contextmanager
 from werkzeug.test import Client, EnvironBuilder
 from flask import _request_ctx_stack
 
 
 class FlaskClient(Client):
-    """Works like a regular Werkzeug test client but has some
-    knowledge about how Flask works to defer the cleanup of the
-    request context stack to the end of a with body when used
-    in a with statement.
+    """Works like a regular Werkzeug test client but has some knowledge about
+    how Flask works to defer the cleanup of the request context stack to the
+    end of a with body when used in a with statement.  For general information
+    about how to use this class refer to :class:`werkzeug.test.Client`.
+
+    Basic usage is outlined in the :ref:`testing` chapter.
     """
 
     preserve_context = context_preserved = False
 
+    @contextmanager
+    def session_transaction(self, *args, **kwargs):
+        """When used in combination with a with statement this opens a
+        session transaction.  This can be used to modify the session that
+        the test client uses.  Once the with block is left the session is
+        stored back.
+
+            with client.session_transaction() as session:
+                session['value'] = 42
+
+        Internally this is implemented by going through a temporary test
+        request context and since session handling could depend on
+        request variables this function accepts the same arguments as
+        :meth:`~flask.Flask.test_request_context` which are directly
+        passed through.
+        """
+        app = self.application
+        environ_overrides = kwargs.pop('environ_overrides', {})
+        if self.cookie_jar is not None:
+            self.cookie_jar.inject_wsgi(environ_overrides)
+        outer_reqctx = _request_ctx_stack.top
+        with app.test_request_context(*args, **kwargs) as c:
+            sess = app.open_session(c.request)
+            if sess is None:
+                raise RuntimeError('Session backend did not open a session. '
+                                   'Check the configuration')
+
+            # Since we have to open a new request context for the session
+            # handling we want to make sure that we hide out own context
+            # from the caller.  By pushing the original request context
+            # (or None) on top of this and popping it we get exactly that
+            # behavior.  It's important to not use the push and pop
+            # methods of the actual request context object since that would
+            # mean that cleanup handlers are called
+            _request_ctx_stack.push(outer_reqctx)
+            try:
+                yield sess
+            finally:
+                _request_ctx_stack.pop()
+
+            resp = app.response_class()
+            if not app.session_interface.is_null_session(sess):
+                app.save_session(sess, resp)
+            if self.cookie_jar is not None:
+                headers = resp.get_wsgi_headers(c.request.environ)
+                self.cookie_jar.extract_wsgi(c.request.environ, headers)
+
     def open(self, *args, **kwargs):
         if self.context_preserved:
             _request_ctx_stack.pop()

tests/flask_tests.py

@@ -1028,6 +1028,50 @@ class BasicFunctionalityTestCase(unittest.TestCase):
             self.assertEqual(rv.data, 'success')
 
 
+class TestToolsTestCase(unittest.TestCase):
+
+    def test_session_transactions(self):
+        app = flask.Flask(__name__)
+        app.testing = True
+        app.secret_key = 'testing'
+
+        @app.route('/')
+        def index():
+            return unicode(flask.session['foo'])
+
+        with app.test_client() as c:
+            with c.session_transaction() as sess:
+                self.assertEqual(len(sess), 0)
+                sess['foo'] = [42]
+                self.assertEqual(len(sess), 1)
+            rv = c.get('/')
+            self.assertEqual(rv.data, '[42]')
+
+    def test_session_transactions_no_null_sessions(self):
+        app = flask.Flask(__name__)
+        app.testing = True
+
+        with app.test_client() as c:
+            try:
+                with c.session_transaction() as sess:
+                    pass
+            except RuntimeError, e:
+                self.assert_('Session backend did not open a session' in str(e))
+            else:
+                self.fail('Expected runtime error')
+
+    def test_session_transactions_keep_context(self):
+        app = flask.Flask(__name__)
+        app.testing = True
+        app.secret_key = 'testing'
+
+        with app.test_client() as c:
+            rv = c.get('/')
+            req = flask.request._get_current_object()
+            with c.session_transaction():
+                self.assert_(req is flask.request._get_current_object())
+
+
 class InstanceTestCase(unittest.TestCase):
 
     def test_explicit_instance_paths(self):
@@ -2209,6 +2253,7 @@ def suite():
     suite.addTest(unittest.makeSuite(SubdomainTestCase))
     suite.addTest(unittest.makeSuite(ViewTestCase))
     suite.addTest(unittest.makeSuite(DeprecationsTestCase))
+    suite.addTest(unittest.makeSuite(TestToolsTestCase))
     suite.addTest(unittest.makeSuite(InstanceTestCase))
     if flask.json_available:
         suite.addTest(unittest.makeSuite(JSONTestCase))