Test print_tokens() with encoding

diff --git a/tests/test_print_tokens.py b/tests/test_print_tokens.py
index af9c9aec80b80b763952deb9a929034990e76c2d..133bb92b40e2ee62abb111f835e61f1a2b1307b7 100644 (file)
--- a/tests/test_print_tokens.py
+++ b/tests/test_print_tokens.py
@@ -1,12 +1,13 @@
 # -*- coding: utf-8 -*-
 
 try:
-    from cStringIO import StringIO
+    from io import StringIO, BytesIO
 except ImportError:
     try:
-        from StringIO import StringIO
+        from cStringIO import StringIO
     except ImportError:
-        from io import StringIO
+        from StringIO import StringIO
+    BytesIO = StringIO
 
 from sqlparse import parse
 from sqlconvert.print_tokens import print_tokens
@@ -17,3 +18,10 @@ def test_print_tokens():
     parsed = parse("select * from T")[0]
     print_tokens(parsed, outfile=sio)
     assert sio.getvalue() == u"SELECT * FROM T"
+
+
+def test_print_tokens_encoded():
+    sio = BytesIO()
+    parsed = parse("select * from T")[0]
+    print_tokens(parsed, outfile=sio, encoding='ascii')
+    assert sio.getvalue() == b"SELECT * FROM T"