Context preserving is now part of Flask and not the test client. This fixes #326

CHANGES

@@ -46,6 +46,9 @@ Relase date to be decided, codename to be chosen.
 - HEAD requests to a method view now automatically dispatch to the `get`
   method if no handler was implemented.
 - Implemented the virtual :mod:`flask.ext` package to import extensions from.
+- The context preservation on exceptions is now an integral component of
+  Flask itself and no longer of the test client.  This cleaned up some
+  internal logic and lowers the odds of runaway request contexts in unittests.
 
 Version 0.7.3
 -------------

flask/ctx.py

@@ -89,6 +89,10 @@ class RequestContext(object):
         self.flashes = None
         self.session = None
 
+        # indicator if the context was preserved.  Next time another context
+        # is pushed the preserved context is popped.
+        self.preserved = False
+
         self.match_request()
 
         # XXX: Support for deprecated functionality.  This is doing away with
@@ -114,6 +118,18 @@ class RequestContext(object):
 
     def push(self):
         """Binds the request context to the current context."""
+        # If an exception ocurrs in debug mode or if context preservation is
+        # activated under exception situations exactly one context stays
+        # on the stack.  The rationale is that you want to access that
+        # information under debug situations.  However if someone forgets to
+        # pop that context again we want to make sure that on the next push
+        # it's invalidated otherwise we run at risk that something leaks
+        # memory.  This is usually only a problem in testsuite since this
+        # functionality is not active in production environments.
+        top = _request_ctx_stack.top
+        if top is not None and top.preserved:
+            top.pop()
+
         _request_ctx_stack.push(self)
 
         # Open the session at the moment that the request context is
@@ -128,8 +144,11 @@ class RequestContext(object):
         also trigger the execution of functions registered by the
         :meth:`~flask.Flask.teardown_request` decorator.
         """
+        self.preserved = False
         self.app.do_teardown_request()
-        _request_ctx_stack.pop()
+        rv = _request_ctx_stack.pop()
+        assert rv is self, 'Popped wrong request context.  (%r instead of %r)' \
+            % (rv, self)
 
     def __enter__(self):
         self.push()
@@ -141,6 +160,16 @@ class RequestContext(object):
         # access the request object in the interactive shell.  Furthermore
         # the context can be force kept alive for the test client.
         # See flask.testing for how this works.
-        if not self.request.environ.get('flask._preserve_context') and \
-           (tb is None or not self.app.preserve_context_on_exception):
+        if self.request.environ.get('flask._preserve_context') or \
+           (tb is not None and self.app.preserve_context_on_exception):
+            self.preserved = True
+        else:
             self.pop()
+
+    def __repr__(self):
+        return '<%s \'%s\' [%s] of %s>' % (
+            self.__class__.__name__,
+            self.request.url,
+            self.request.method,
+            self.app.name
+        )

flask/globals.py

@@ -19,6 +19,7 @@ def _lookup_object(name):
         raise RuntimeError('working outside of request context')
     return getattr(top, name)
 
+
 # context locals
 _request_ctx_stack = LocalStack()
 current_app = LocalProxy(partial(_lookup_object, 'app'))

flask/testing.py

@@ -37,7 +37,7 @@ class FlaskClient(Client):
     Basic usage is outlined in the :ref:`testing` chapter.
     """
 
-    preserve_context = context_preserved = False
+    preserve_context = False
 
     @contextmanager
     def session_transaction(self, *args, **kwargs):
@@ -88,7 +88,6 @@ class FlaskClient(Client):
             self.cookie_jar.extract_wsgi(c.request.environ, headers)
 
     def open(self, *args, **kwargs):
-        self._pop_reqctx_if_necessary()
         kwargs.setdefault('environ_overrides', {}) \
             ['flask._preserve_context'] = self.preserve_context
 
@@ -97,14 +96,10 @@ class FlaskClient(Client):
         follow_redirects = kwargs.pop('follow_redirects', False)
         builder = make_test_environ_builder(self.application, *args, **kwargs)
 
-        old = _request_ctx_stack.top
-        try:
-            return Client.open(self, builder,
-                               as_tuple=as_tuple,
-                               buffered=buffered,
-                               follow_redirects=follow_redirects)
-        finally:
-            self.context_preserved = _request_ctx_stack.top is not old
+        return Client.open(self, builder,
+                           as_tuple=as_tuple,
+                           buffered=buffered,
+                           follow_redirects=follow_redirects)
 
     def __enter__(self):
         if self.preserve_context:
@@ -114,12 +109,10 @@ class FlaskClient(Client):
 
     def __exit__(self, exc_type, exc_value, tb):
         self.preserve_context = False
-        self._pop_reqctx_if_necessary()
-
-    def _pop_reqctx_if_necessary(self):
-        if self.context_preserved:
-            # we have to use _request_ctx_stack.top.pop instead of
-            # _request_ctx_stack.pop since we want teardown handlers
-            # to be executed.
-            _request_ctx_stack.top.pop()
-            self.context_preserved = False
+
+        # on exit we want to clean up earlier.  Normally the request context
+        # stays preserved until the next request in the same thread comes
+        # in.  See RequestGlobals.push() for the general behavior.
+        top = _request_ctx_stack.top
+        if top is not None and top.preserved:
+            top.pop()

flask/testsuite/basic.py

@@ -904,6 +904,23 @@ class BasicFunctionalityTestCase(FlaskTestCase):
         self.assertEqual(c.get('/bar/').data, 'bar')
         self.assertEqual(c.get('/bar/123').data, '123')
 
+    def test_preserve_only_once(self):
+        app = flask.Flask(__name__)
+        app.debug = True
+
+        @app.route('/fail')
+        def fail_func():
+            1/0
+
+        c = app.test_client()
+        for x in xrange(3):
+            with self.assert_raises(ZeroDivisionError):
+                c.get('/fail')
+
+        self.assert_(flask._request_ctx_stack.top is not None)
+        flask._request_ctx_stack.pop()
+        self.assert_(flask._request_ctx_stack.top is None)
+
 
 class ContextTestCase(FlaskTestCase):