1 files changed, 7 insertions, 2 deletions
diff --git a/morphlib/bins_tests.py b/morphlib/bins_tests.py
index 3895680f..e2bb2f6a 100644
--- a/morphlib/bins_tests.py
+++ b/morphlib/bins_tests.py
@@ -17,13 +17,18 @@ import gzip
 import os
 import shutil
 import stat
+import sys
 import tempfile
 import tarfile
 import unittest
-import StringIO
 
 import morphlib
 
+if sys.version_info[0] == 2:
+    from StringIO import StringIO
+else:
+    from io import StringIO
+
 
 class BinsTest(unittest.TestCase):
 
@@ -152,7 +157,7 @@ class ExtractTests(unittest.TestCase):
         shutil.rmtree(self.tempdir)
 
     def create_chunk(self, callback):
-        fh = StringIO.StringIO()
+        fh = StringIO()
         os.mkdir(self.instdir)
         patterns = callback(self.instdir)
         morphlib.bins.create_chunk(self.instdir, fh, patterns)