Fix CSV exports in Python 3 by converting StringIO to BytesIO (#1107)

* Fix CSV exports in Python 3 by converting StringIO to BytesIO
bulk-clear-sessions
Kevin Chung 2019-09-10 03:22:01 -04:00 committed by GitHub
parent 607c517d28
commit b8c1970b8e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 25 additions and 2 deletions

View File

@ -125,8 +125,8 @@ def export_csv():
if model is None:
abort(404)
output = six.StringIO()
writer = csv.writer(output)
temp = six.StringIO()
writer = csv.writer(temp)
header = [column.name for column in model.__mapper__.columns]
writer.writerow(header)
@ -138,7 +138,14 @@ def export_csv():
[getattr(curr, column.name) for column in model.__mapper__.columns]
)
temp.seek(0)
# In Python 3 send_file requires bytes
output = six.BytesIO()
output.write(temp.getvalue().encode("utf-8"))
output.seek(0)
temp.close()
return send_file(
output,
as_attachment=True,

View File

@ -0,0 +1,16 @@
from tests.helpers import create_ctfd, destroy_ctfd, login_as_user, gen_challenge
def test_export_csv_works():
"""Test that CSV exports work properly"""
app = create_ctfd()
with app.app_context():
gen_challenge(app.db)
client = login_as_user(app, name="admin", password="password")
csv_data = client.get("/admin/export/csv?table=challenges").get_data(
as_text=True
)
assert len(csv_data) > 0
destroy_ctfd(app)