Updated tests to be py3 compatible

ParallelSSH · Feb 21, 2017 · f64ad59 · f64ad59
1 parent 4256d74
commit f64ad59
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/tests/test_ssh_client.py b/tests/test_ssh_client.py
@@ -37,6 +37,12 @@
 import random, string
 import tempfile
 
+
+try:
+    xrange
+except NameError:
+    xrange = range
+
 USER_KEY_PATH = os.path.sep.join([os.path.dirname(__file__), 'test_client_private_key'])
 USER_KEY = paramiko.RSAKey.from_private_key_file(USER_KEY_PATH)
 

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -1,7 +1,10 @@
 from pssh import utils
 import unittest
 import os
-from cStringIO import StringIO
+try:
+    from cStringIO import StringIO as BytesIO
+except ImportError:
+    from io import BytesIO
 from uuid import uuid4
 
 PKEY_FILENAME = os.path.sep.join([os.path.dirname(__file__), 'test_client_private_key'])
@@ -14,19 +17,19 @@ def test_enabling_host_logger(self):
         utils.enable_host_logger()
         # And again to test only one handler is attached
         utils.enable_host_logger()
-        self.assertTrue(len(utils.host_logger.handlers)==1)
+        self.assertTrue(len(utils.host_logger.handlers) == 1)
 
     def test_enabling_pssh_logger(self):
         utils.enable_logger(utils.logger)
-        self.assertTrue(len(utils.logger.handlers)==1)
+        self.assertTrue(len(utils.logger.handlers) == 1)
 
     def test_loading_key_files(self):
         for key_filename in [PKEY_FILENAME, DSA_KEY_FILENAME, ECDSA_KEY_FILENAME]:
             pkey = utils.load_private_key(key_filename)
             self.assertTrue(pkey, msg="Error loading key from file %s" % (key_filename,))
             pkey = utils.load_private_key(open(key_filename))
             self.assertTrue(pkey, msg="Error loading key from open file object for file %s" % (key_filename,))
-        fake_key = StringIO("blah blah fakey fakey key")
+        fake_key = BytesIO(b"blah blah fakey fakey key")
         self.assertFalse(utils.load_private_key(fake_key))
 
     def test_openssh_config_missing(self):