Added flask.stream_with_context

flask/testsuite/appctx.py

@@ -75,6 +75,26 @@ class AppContextTestCase(FlaskTestCase):
             self.assert_equal(
                 flask.render_template_string('{{ g.spam }}'), 'eggs')
 
+    def test_context_refcounts(self):
+        called = []
+        app = flask.Flask(__name__)
+        @app.teardown_request
+        def teardown_req(error=None):
+            called.append('request')
+        @app.teardown_appcontext
+        def teardown_app(error=None):
+            called.append('app')
+        @app.route('/')
+        def index():
+            with flask._app_ctx_stack.top:
+                with flask._request_ctx_stack.top:
+                    pass
+            self.assert_(flask._request_ctx_stack.request.environ
+                ['werkzeug.request'] is not None)
+        c = app.test_client()
+        c.get('/')
+        self.assertEqual(called, ['request', 'app'])
+
 
 def suite():
     suite = unittest.TestSuite()

flask/testsuite/helpers.py

@@ -397,6 +397,64 @@ class NoImportsTestCase(FlaskTestCase):
             self.fail('Flask(import_name) is importing import_name.')
 
 
+class StreamingTestCase(FlaskTestCase):
+
+    def test_streaming_with_context(self):
+        app = flask.Flask(__name__)
+        app.testing = True
+        @app.route('/')
+        def index():
+            def generate():
+                yield 'Hello '
+                yield flask.request.args['name']
+                yield '!'
+            return flask.Response(flask.stream_with_context(generate()))
+        c = app.test_client()
+        rv = c.get('/?name=World')
+        self.assertEqual(rv.data, 'Hello World!')
+
+    def test_streaming_with_context_as_decorator(self):
+        app = flask.Flask(__name__)
+        app.testing = True
+        @app.route('/')
+        def index():
+            @flask.stream_with_context
+            def generate():
+                yield 'Hello '
+                yield flask.request.args['name']
+                yield '!'
+            return flask.Response(generate())
+        c = app.test_client()
+        rv = c.get('/?name=World')
+        self.assertEqual(rv.data, 'Hello World!')
+
+    def test_streaming_with_context_and_custom_close(self):
+        app = flask.Flask(__name__)
+        app.testing = True
+        called = []
+        class Wrapper(object):
+            def __init__(self, gen):
+                self._gen = gen
+            def __iter__(self):
+                return self
+            def close(self):
+                called.append(42)
+            def next(self):
+                return self._gen.next()
+        @app.route('/')
+        def index():
+            def generate():
+                yield 'Hello '
+                yield flask.request.args['name']
+                yield '!'
+            return flask.Response(flask.stream_with_context(
+                Wrapper(generate())))
+        c = app.test_client()
+        rv = c.get('/?name=World')
+        self.assertEqual(rv.data, 'Hello World!')
+        self.assertEqual(called, [42])
+
+
 def suite():
     suite = unittest.TestSuite()
     if flask.json_available:
@@ -404,4 +462,5 @@ def suite():
     suite.addTest(unittest.makeSuite(SendfileTestCase))
     suite.addTest(unittest.makeSuite(LoggingTestCase))
     suite.addTest(unittest.makeSuite(NoImportsTestCase))
+    suite.addTest(unittest.makeSuite(StreamingTestCase))
     return suite