Merge pull request pallets-eco#2 from pawl/test_csv_export_3

Add tests for csv export and fix python 3 unicode issues

flask_admin/model/base.py

@@ -22,7 +22,7 @@
                                  get_redirect_target, flash_errors)
 from flask_admin.tools import rec_getattr
 from flask_admin._backwards import ObsoleteAttr
-from flask_admin._compat import iteritems, OrderedDict, as_unicode
+from flask_admin._compat import iteritems, OrderedDict, as_unicode, csv_encode
 from .helpers import prettify_name, get_mdict_item_or_list
 from .ajax import AjaxModelLoader
 from .fields import ListEditableFieldList
@@ -1942,25 +1942,19 @@ def write(self, value):
         writer = csv.writer(Echo())
 
         def generate():
-            # Needed as python 2 csvwriter does not support unicode
-            def fix_unicode(t):
-                return as_unicode(t).encode('utf-8')
-
             # Append the column titles at the beginning
-            titles = [fix_unicode(c[1]) for c in self._list_columns]
+            titles = [csv_encode(c[1]) for c in self._list_columns]
             yield writer.writerow(titles)
 
             for row in data:
-                vals = [fix_unicode(self.get_export_value(row, c[0]))
+                vals = [csv_encode(self.get_export_value(row, c[0]))
                         for c in self._list_columns]
                 yield writer.writerow(vals)
 
-        filename = '{}_{}.csv'.format(
-            self.name,
-            time.strftime("%Y-%m-%d_%H-%M-%S")
-        )
+        filename = '%s_%s.csv' % (self.name,
+                                  time.strftime("%Y-%m-%d_%H-%M-%S"))
 
-        disposition = 'attachment;filename={}'.format(secure_filename(filename))
+        disposition = 'attachment;filename=%s' % (secure_filename(filename),)
 
         return Response(
             stream_with_context(generate()),

flask_admin/tests/test_model.py

@@ -12,6 +12,7 @@
 from flask_admin import Admin, form
 from flask_admin._compat import iteritems, itervalues
 from flask_admin.model import base, filters
+from itertools import islice
 
 
 def wtforms2_and_up(func):
@@ -46,8 +47,8 @@ def operation(self):
 
 
 class MockModelView(base.BaseModelView):
-    def __init__(self, model, name=None, category=None, endpoint=None, url=None,
-                 **kwargs):
+    def __init__(self, model, data=None, name=None, category=None,
+                 endpoint=None, url=None, **kwargs):
         # Allow to set any attributes from parameters
         for k, v in iteritems(kwargs):
             setattr(self, k, v)
@@ -60,9 +61,12 @@ def __init__(self, model, name=None, category=None, endpoint=None, url=None,
 
         self.search_arguments = []
 
-        self.all_models = {1: Model(1),
-                           2: Model(2)}
-        self.last_id = 3
+        if data is None:
+            self.all_models = {1: Model(1), 2: Model(2)}
+        else:
+            self.all_models = data
+
+        self.last_id = len(self.all_models) + 1
 
     # Scaffolding
     def get_pk_value(self, model):
@@ -89,9 +93,12 @@ def scaffold_form(self):
         return Form
 
     # Data
-    def get_list(self, page, sort_field, sort_desc, search, filters):
+    def get_list(self, page, sort_field, sort_desc, search, filters,
+                 page_size=None):
         self.search_arguments.append((page, sort_field, sort_desc, search, filters))
-        return len(self.all_models), itervalues(self.all_models)
+        count = len(self.all_models)
+        data = islice(itervalues(self.all_models), 0, page_size)
+        return count, data
 
     def get_one(self, id):
         return self.all_models.get(int(id))
@@ -538,3 +545,88 @@ class DummyView(MockModelView):
 
     view = DummyView(Model)
     eq_(view.name, 'Dummy View')
+
+
+def test_export_csv():
+    app, admin = setup()
+    client = app.test_client()
+
+    # test redirect when csv export is disabled
+    view = MockModelView(Model, column_list=['col1', 'col2'], endpoint="test")
+    admin.add_view(view)
+
+    rv = client.get('/admin/test/export/csv/')
+    eq_(rv.status_code, 302)
+
+    # basic test of csv export with a few records
+    view_data = {
+        1: Model(1, "col1_1", "col2_1"),
+        2: Model(2, "col1_2", "col2_2"),
+        3: Model(3, "col1_3", "col2_3"),
+    }
+
+    view = MockModelView(Model, view_data, can_export=True,
+                         column_list=['col1', 'col2'])
+    admin.add_view(view)
+
+    rv = client.get('/admin/model/export/csv/')
+    data = rv.data.decode('utf-8')
+    eq_(rv.mimetype, 'text/csv')
+    eq_(rv.status_code, 200)
+    ok_("Col1,Col2\r\n"
+        "col1_1,col2_1\r\n"
+        "col1_2,col2_2\r\n"
+        "col1_3,col2_3\r\n" == data)
+
+    # test utf8 characters in csv export
+    view_data[4] = Model(1, u'\u2013ut8_1\u2013', u'\u2013utf8_2\u2013')
+    view = MockModelView(Model, view_data, can_export=True,
+                         column_list=['col1', 'col2'], endpoint="utf8")
+    admin.add_view(view)
+
+    rv = client.get('/admin/utf8/export/csv/')
+    data = rv.data.decode('utf-8')
+    eq_(rv.status_code, 200)
+    ok_(u'\u2013ut8_1\u2013,\u2013utf8_2\u2013\r\n' in data)
+
+    # test row limit
+    view_data = {
+        1: Model(1, "col1_1", "col2_1"),
+        2: Model(2, "col1_2", "col2_2"),
+        3: Model(3, "col1_3", "col2_3"),
+    }
+
+    view = MockModelView(Model, view_data, can_export=True,
+                         column_list=['col1', 'col2'], export_max_rows=2,
+                         endpoint='row_limit_2')
+    admin.add_view(view)
+
+    rv = client.get('/admin/row_limit_2/export/csv/')
+    data = rv.data.decode('utf-8')
+    eq_(rv.status_code, 200)
+    ok_("Col1,Col2\r\n"
+        "col1_1,col2_1\r\n"
+        "col1_2,col2_2\r\n" == data)
+
+    # test None type, integer type, column_labels, and column_formatters
+    view_data = {
+        1: Model(1, "col1_1", 1),
+        2: Model(2, "col1_2", 2),
+        3: Model(3, None, 3),
+    }
+
+    view = MockModelView(
+        Model, view_data, can_export=True, column_list=['col1', 'col2'],
+        column_labels={'col1': 'Str Field', 'col2': 'Int Field'},
+        column_formatters=dict(col2=lambda v, c, m, p: m.col2*2),
+        endpoint="types_and_formatters"
+    )
+    admin.add_view(view)
+
+    rv = client.get('/admin/types_and_formatters/export/csv/')
+    data = rv.data.decode('utf-8')
+    eq_(rv.status_code, 200)
+    ok_("Str Field,Int Field\r\n"
+        "col1_1,2\r\n"
+        "col1_2,4\r\n"
+        ",6\r\n" == data)