test: import StringIO from the io module when on python3

xtuml · Nov 4, 2017 · 897ea4d · 897ea4d
1 parent b03bfea
commit 897ea4d
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -5,9 +5,15 @@
 import tempfile
 import sys
 import os
-import StringIO
 import logging
 
+try:
+    # python2
+    from StringIO import StringIO
+except ImportError:
+    # python3
+    from io import StringIO
+
 import rsl
 
 
@@ -16,7 +22,7 @@ class TestCommandLineInterface(unittest.TestCase):
     def setUp(self):
         self.temp_files = set()
         logger = logging.getLogger()
-        self.log = logging.StreamHandler(StringIO.StringIO())
+        self.log = logging.StreamHandler(StringIO())
         logger.addHandler(self.log)
 
     def tearDown(self):