X-Git-Url: https://git.phdru.name/?a=blobdiff_plain;f=tests%2Ftest_print_tokens.py;h=133bb92b40e2ee62abb111f835e61f1a2b1307b7;hb=433f46b3a387ec5d4daf2c057923c6d4c1bc263a;hp=af9c9aec80b80b763952deb9a929034990e76c2d;hpb=6978b242d0e02f65473930e9e2a0f2665c2e8863;p=sqlconvert.git

diff --git a/tests/test_print_tokens.py b/tests/test_print_tokens.py
index af9c9ae..133bb92 100644
--- 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"