瀏覽代碼

HUE-8737 [core] Fix desktop lib py3 issues
desktop.lib.django_util_test
desktop.lib.test_export_csvxls

Ying Chen 6 年之前
父節點
當前提交
dcd3303cc7

+ 4 - 2
desktop/core/src/desktop/lib/django_util.py

@@ -254,7 +254,7 @@ def render_injected(http_resp, extra_html):
     return http_resp
 
   # Look for the </body> tag and inject the popup <div>
-  markers = ('</body>', '</BODY>')
+  markers = (b'</body>', b'</BODY>')
   content = http_resp.content
   for marker in markers:
     pos = content.rfind(marker)
@@ -268,7 +268,9 @@ def render_injected(http_resp, extra_html):
     extra_html = extra_html.html
   if callable(extra_html):
     extra_html = extra_html()
-  http_resp.content = ''.join((content[:pos], extra_html, content[pos:]))
+  if not isinstance(extra_html, bytes):
+    extra_html = extra_html.encode('utf-8')
+  http_resp.content = b''.join((content[:pos], extra_html, content[pos:]))
   return http_resp
 
 

+ 3 - 3
desktop/core/src/desktop/lib/django_util_test.py

@@ -151,15 +151,15 @@ def test_popup_injection():
   """Test that result injection works"""
   base = HttpResponse('<html><head></head><body>Hello</body></html>')
   resp = django_util.render_injected(base, ' Cookie monster')
-  assert_true('Hello Cookie monster' in resp.content)
+  assert_true(b'Hello Cookie monster' in resp.content)
 
   redirect = HttpResponseRedirect('http://www.cnn.com')
   resp = django_util.render_injected(redirect, 'Cookie monster')
-  assert_true('Cookie monster' not in resp.content)
+  assert_true(b'Cookie monster' not in resp.content)
 
   json = django_util.render_json('blah')
   resp = django_util.render_injected(json, 'Cookie monster')
-  assert_true('Cookie monster' not in resp.content)
+  assert_true(b'Cookie monster' not in resp.content)
 
   assert_raises(AssertionError, django_util.render_injected, "foo", "bar")
 

+ 3 - 2
desktop/core/src/desktop/lib/export_csvxls.py

@@ -37,7 +37,7 @@ from django.utils.http import urlquote
 from desktop.lib import i18n
 
 if sys.version_info[0] > 2:
-  from io import StringIO as string_io
+  from io import BytesIO as string_io
 else:
   from StringIO import StringIO as string_io
 
@@ -154,7 +154,8 @@ def make_response(generator, format, name, encoding=None, user_agent=None): #TOD
 
   try:
     name = name.encode('ascii')
-    resp['Content-Disposition'] = 'attachment; filename="%s.%s"' % (name, format)
+    format = format.encode('ascii')
+    resp['Content-Disposition'] = b'attachment; filename="%s.%s"' % (name, format)
   except UnicodeEncodeError:
     name = urlquote(name)
     if user_agent is not None and 'Firefox' in user_agent:

+ 8 - 8
desktop/core/src/desktop/lib/test_export_csvxls.py

@@ -26,7 +26,7 @@ from openpyxl import load_workbook
 from desktop.lib.export_csvxls import create_generator, make_response
 
 if sys.version_info[0] > 2:
-  from io import StringIO as string_io
+  from io import BytesIO as string_io
 else:
   from cStringIO import StringIO as string_io
 
@@ -42,16 +42,16 @@ def test_export_csv():
   generator = create_generator(content_generator(headers, data), "csv")
   response = make_response(generator, "csv", "foo")
   assert_equal("application/csv", response["content-type"])
-  content = ''.join(response.streaming_content)
-  assert_equal('x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\nhttp://gethue.com,http://gethue.com\r\n', content)
+  content = b''.join(response.streaming_content)
+  assert_equal(b'x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\nhttp://gethue.com,http://gethue.com\r\n', content)
   assert_equal('attachment; filename="foo.csv"', response["content-disposition"])
 
   # Check non-ASCII for any browser except FF or no browser info
   generator = create_generator(content_generator(headers, data), "csv")
   response = make_response(generator, "csv", u'gんtbhんjk?¥n')
   assert_equal("application/csv", response["content-type"])
-  content = ''.join(response.streaming_content)
-  assert_equal('x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\nhttp://gethue.com,http://gethue.com\r\n', content)
+  content = b''.join(response.streaming_content)
+  assert_equal(b'x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\nhttp://gethue.com,http://gethue.com\r\n', content)
   assert_equal('attachment; filename="g%E3%82%93tbh%E3%82%93jk%EF%BC%9F%EF%BF%A5n.csv"', response["content-disposition"])
 
   # Check non-ASCII for FF browser
@@ -59,8 +59,8 @@ def test_export_csv():
   response = make_response(generator, "csv", u'gんtbhんjk?¥n',
                            user_agent='Mozilla / 5.0(Macintosh; Intel Mac OS X 10.12;rv:59.0) Gecko / 20100101 Firefox / 59.0)')
   assert_equal("application/csv", response["content-type"])
-  content = ''.join(response.streaming_content)
-  assert_equal('x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\nhttp://gethue.com,http://gethue.com\r\n', content)
+  content = b''.join(response.streaming_content)
+  assert_equal(b'x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\nhttp://gethue.com,http://gethue.com\r\n', content)
   assert_equal('attachment; filename*="g%E3%82%93tbh%E3%82%93jk%EF%BC%9F%EF%BF%A5n.csv"',
                response["content-disposition"])
 
@@ -84,7 +84,7 @@ def test_export_xls():
 
 
 def _read_xls_sheet_data(response):
-  content = ''.join(response.content)
+  content = bytes(response.content)
 
   data = string_io()
   data.write(content)