test flask streaming responses

ckan · May 25, 2017 · 2d30527 · 2d30527
1 parent f24fb3d
commit 2d30527
Show file tree

Hide file tree

Showing 8 changed files with 193 additions and 8 deletions.
diff --git a/ckan/common.py b/ckan/common.py
@@ -41,18 +41,21 @@ def is_flask_request():
              not pylons_request_available))
 
 
-def streaming_response(data):
+def streaming_response(
+        data, mimetype='application/octet-stream', with_context=False):
     iter_data = iter(data)
     if is_flask_request():
-        # Removal of context variables for pylon's app prevented
-        # inside `pylons_app.py`, but flask requires special treating.
-        # Otherwice we are going to constantly receive errors about
-        # usage of unregistered values from context.
-        resp = flask.Response(
-            flask.stream_with_context(iter_data))
+        # Removal of context variables for pylon's app is prevented
+        # inside `pylons_app.py`. It would be better to decide on the fly
+        # whether we need to preserve context, but it won't affect performance
+        # in any visible way and we are going to get rid of pylons anyway.
+        # Flask allows to do this in easy way.
+        if with_context:
+            iter_data = flask.stream_with_context(iter_data)
+        resp = flask.Response(iter_data, mimetype=mimetype)
     else:
         response.app_iter = iter_data
-        resp = response
+        resp = response.headers['Content-type'] = mimetype
     return resp
 
 

diff --git a/ckanext/example_flask_streaming/__init__.py b/ckanext/example_flask_streaming/__init__.py
diff --git a/ckanext/example_flask_streaming/plugin.py b/ckanext/example_flask_streaming/plugin.py
@@ -0,0 +1,87 @@
+# encoding: utf-8
+
+import os.path as path
+
+from flask import Blueprint
+import flask
+
+import ckan.plugins as p
+from ckan.common import streaming_response
+
+
+def stream_string():
+    u'''A simple view function'''
+    def generate():
+        for w in u'Hello World, this is served from an extension'.split():
+            yield w
+    return streaming_response(generate())
+
+
+def stream_template(**kwargs):
+    u'''A simple replacement for the pylons About page.'''
+    tpl = flask.current_app.jinja_env.get_template('stream.html')
+    gen = tpl.stream(kwargs)
+    gen.enable_buffering()
+    return streaming_response(gen)
+
+
+def stream_file():
+    u'''A simple replacement for the flash Hello view function.'''
+    f_path = path.join(
+        path.dirname(path.abspath(__file__)), 'tests/10lines.txt')
+
+    def gen():
+        with open(f_path) as test_file:
+            for line in test_file:
+                yield line
+
+    return streaming_response(gen())
+
+
+def stream_context():
+    u'''A simple replacement for the flash Hello view function.'''
+    html = '''{{ request.args.var }}'''
+
+    def gen():
+        yield flask.render_template_string(html)
+
+    return streaming_response(gen(), with_context=True)
+
+
+def stream_without_context():
+    u'''A simple replacement for the flash Hello view function.'''
+    html = '''{{ request.args.var }}'''
+
+    def gen():
+        yield flask.render_template_string(html)
+
+    return streaming_response(gen())
+
+
+class ExampleFlaskStreamingPlugin(p.SingletonPlugin):
+    u'''
+    An example plugin to demonstrate Flask streaming responses.
+    '''
+    p.implements(p.IBlueprint)
+
+    def get_blueprint(self):
+        u'''Return a Flask Blueprint object to be registered by the app.'''
+
+        # Create Blueprint for plugin
+        blueprint = Blueprint(self.name, self.__module__)
+        blueprint.template_folder = u'templates'
+        # Add plugin url rules to Blueprint object
+        rules = [
+            (u'/stream/string', u'stream_string', stream_string),
+            (u'/stream/template/<int:count>', u'stream_template',
+             stream_template),
+            (u'/stream/template/', u'stream_template', stream_template),
+            (u'/stream/file', u'stream_file', stream_file),
+            (u'/stream/context', u'stream_context', stream_context),
+            (u'/stream/without_context', u'stream_without_context',
+             stream_without_context),
+        ]
+        for rule in rules:
+            blueprint.add_url_rule(*rule)
+
+        return blueprint
diff --git a/ckanext/example_flask_streaming/templates/stream.html b/ckanext/example_flask_streaming/templates/stream.html
@@ -0,0 +1,13 @@
+<!DOCTYPE html>
+<html>
+<head>
+    <title>My New stream Page</title>
+</head>
+<body>
+    <h1>This is an stream page served from an extention.</h1>
+    {% for i in range(count) %}
+      <p>{{ i }}</p>
+    {% endfor %}
+
+</body>
+</html>
diff --git a/ckanext/example_flask_streaming/tests/10lines.txt b/ckanext/example_flask_streaming/tests/10lines.txt
@@ -0,0 +1,10 @@
+Vel quam elementum pulvinar etiam. Semper viverra nam libero justo, laoreet sit amet cursus sit amet, dictum sit amet justo donec enim diam, vulputate ut pharetra sit amet, aliquam id.
+In massa tempor nec feugiat nisl pretium fusce id velit! Odio pellentesque diam volutpat commodo sed egestas egestas fringilla phasellus faucibus scelerisque eleifend donec pretium vulputate sapien nec sagittis aliquam?
+Bibendum neque egestas congue quisque egestas diam in. Rhoncus est pellentesque elit ullamcorper dignissim cras tincidunt lobortis feugiat vivamus at augue eget arcu dictum varius duis at consectetur lorem donec!
+Aliquet bibendum enim, facilisis gravida neque convallis a cras semper auctor neque, vitae tempus quam pellentesque nec nam aliquam sem et tortor consequat id porta! A cras semper auctor neque?
+Amet justo donec enim diam, vulputate ut pharetra sit amet, aliquam. Non quam lacus suspendisse faucibus interdum posuere lorem ipsum dolor sit amet, consectetur adipiscing elit duis tristique sollicitudin nibh!
+Senectus et netus et malesuada fames ac turpis egestas sed tempus, urna et pharetra pharetra, massa! Urna nunc id cursus metus aliquam eleifend mi in nulla posuere sollicitudin aliquam ultrices!
+Amet, dictum sit amet justo donec enim diam, vulputate ut? At urna condimentum mattis pellentesque id nibh tortor, id aliquet lectus proin nibh nisl, condimentum id venenatis a, condimentum vitae?
+Vestibulum, lectus mauris ultrices eros, in cursus turpis massa tincidunt dui ut ornare lectus sit amet est placerat. Imperdiet nulla malesuada pellentesque elit eget gravida cum sociis natoque penatibus et.
+Sed vulputate mi sit amet mauris commodo quis imperdiet massa tincidunt nunc pulvinar sapien et ligula ullamcorper malesuada. Risus pretium quam vulputate dignissim suspendisse in est ante in nibh mauris!
+Eget nullam non nisi est, sit. Aliquet eget sit amet tellus cras adipiscing enim eu turpis egestas pretium aenean pharetra, magna ac placerat vestibulum, lectus mauris ultrices eros, in cursus.
diff --git a/ckanext/example_flask_streaming/tests/__init__.py b/ckanext/example_flask_streaming/tests/__init__.py
diff --git a/ckanext/example_flask_streaming/tests/test_streaming_responses.py b/ckanext/example_flask_streaming/tests/test_streaming_responses.py
@@ -0,0 +1,71 @@
+# encoding: utf-8
+
+import os.path as path
+
+from nose.tools import eq_, assert_raises
+from webtest.app import TestRequest
+from webtest import lint  # NOQA
+import ckan.plugins as plugins
+import ckan.tests.helpers as helpers
+
+
+class TestFlaskStreaming(helpers.FunctionalTestBase):
+
+    def _get_resp(self, url):
+        req = TestRequest.blank(url)
+        app = lint.middleware(self.app.app)
+        res = req.get_response(app, True)
+        return res
+
+    def setup(self):
+        self.app = helpers._get_test_app()
+
+        # Install plugin and register its blueprint
+        if not plugins.plugin_loaded(u'example_flask_streaming'):
+            plugins.load(u'example_flask_streaming')
+            plugin = plugins.get_plugin(u'example_flask_streaming')
+            self.app.flask_app.register_extension_blueprint(
+                plugin.get_blueprint())
+
+    def test_accordance_of_chunks(self):
+        u'''Test extension sets up a unique route.'''
+        url = '/stream/string'
+        resp = self._get_resp(url)
+        eq_(
+            u'Hello World, this is served from an extension'.split(),
+            list(resp.app_iter))
+        resp.app_iter.close()
+
+    def test_template_streaming(self):
+        u'''Test extension sets up a unique route.'''
+        url = '/stream/template'
+        resp = self._get_resp(url)
+        eq_(1, len(list(resp.app_iter)))
+
+        url = '/stream/template/7'
+        resp = self._get_resp(url)
+        eq_(2, len(list(resp.app_iter)))
+        resp._app_iter.close()
+
+    def test_file_streaming(self):
+        u'''Test extension sets up a unique route.'''
+        url = '/stream/file'
+        resp = self._get_resp(url)
+        f_path = path.join(path.dirname(path.abspath(__file__)), '10lines.txt')
+        with open(f_path) as test_file:
+            content = test_file.readlines()
+            eq_(content, list(resp.app_iter))
+        resp._app_iter.close()
+
+    def test_render_with_context(self):
+        u'''Test extension sets up a unique route.'''
+        url = '/stream/context?var=10'
+        resp = self._get_resp(url)
+        eq_('10', resp.body)
+
+    def test_render_without_context(self):
+        u'''Test extension sets up a unique route.'''
+        url = '/stream/without_context?var=10'
+        resp = self._get_resp(url)
+        assert_raises(AttributeError, str.join, '', resp.app_iter)
+        resp.app_iter.close()
diff --git a/setup.py b/setup.py
@@ -159,6 +159,7 @@
         'example_iconfigurer_v1 = ckanext.example_iconfigurer.plugin_v1:ExampleIConfigurerPlugin',
         'example_iconfigurer_v2 = ckanext.example_iconfigurer.plugin_v2:ExampleIConfigurerPlugin',
         'example_flask_iblueprint = ckanext.example_flask_iblueprint.plugin:ExampleFlaskIBlueprintPlugin',
+        'example_flask_streaming = ckanext.example_flask_streaming.plugin:ExampleFlaskStreamingPlugin',
         'example_iuploader = ckanext.example_iuploader.plugin:ExampleIUploader',
         'example_ipermissionlabels = ckanext.example_ipermissionlabels.plugin:ExampleIPermissionLabelsPlugin',
     ],