Update tests to use unittest2.

6 files changed, 30 insertions, 28 deletions

diff --git a/tests/test_bigfile.py b/tests/test_bigfile.py
index 39bd095..02e052e 100644
--- a/tests/test_bigfile.py
+++ b/tests/test_bigfile.py
@@ -4,12 +4,12 @@ try:
     from StringIO import StringIO
 except ImportError:
     from io import StringIO
-import unittest
+import unittest2
 
 import rsa
 from rsa import bigfile, varblock, pkcs1
 
-class BigfileTest(unittest.TestCase):
+class BigfileTest(unittest2.TestCase):
 
     def test_encrypt_decrypt_bigfile(self):
 
diff --git a/tests/test_integers.py b/tests/test_integers.py
index d4fa087..0a712aa 100644
--- a/tests/test_integers.py
+++ b/tests/test_integers.py
@@ -1,10 +1,10 @@
 '''Tests integer operations.'''
 
-import unittest
+import unittest2
 
 import rsa.core
 
-class IntegerTest(unittest.TestCase):
+class IntegerTest(unittest2.TestCase):
 
     def setUp(self):
         (self.pub, self.priv) = rsa.newkeys(64)
diff --git a/tests/test_load_save_keys.py b/tests/test_load_save_keys.py
index 56d45c4..fabe92f 100644
--- a/tests/test_load_save_keys.py
+++ b/tests/test_load_save_keys.py
@@ -1,7 +1,7 @@
 '''Unittest for saving and loading keys.'''
 
 import base64
-import unittest
+import unittest2
 from rsa._compat import b
 
 import rsa.key
@@ -53,7 +53,7 @@ CLEAN_PUBLIC_PEM = b('''\
 ''' % B64PUB_DER)
 
 
-class DerTest(unittest.TestCase):
+class DerTest(unittest2.TestCase):
     '''Test saving and loading DER keys.'''
 
     def test_load_private_key(self):
@@ -88,7 +88,7 @@ class DerTest(unittest.TestCase):
 
         self.assertEqual(PUBLIC_DER, der)
 
-class PemTest(unittest.TestCase):
+class PemTest(unittest2.TestCase):
     '''Test saving and loading PEM keys.'''
 
 
diff --git a/tests/test_pkcs1.py b/tests/test_pkcs1.py
index c841485..d8fb1b4 100644
--- a/tests/test_pkcs1.py
+++ b/tests/test_pkcs1.py
@@ -1,12 +1,12 @@
 '''Tests string operations.'''
 
 import struct
-import unittest
+import unittest2
 
 import rsa
 from rsa import pkcs1
 
-class BinaryTest(unittest.TestCase):
+class BinaryTest(unittest2.TestCase):
 
     def setUp(self):
         (self.pub, self.priv) = rsa.newkeys(256)
@@ -46,7 +46,7 @@ class BinaryTest(unittest.TestCase):
         
         self.assertNotEqual(encrypted1, encrypted2)
 
-class SignatureTest(unittest.TestCase):
+class SignatureTest(unittest2.TestCase):
 
     def setUp(self):
         (self.pub, self.priv) = rsa.newkeys(512)
diff --git a/tests/test_strings.py b/tests/test_strings.py
index 58d3833..001456d 100644
--- a/tests/test_strings.py
+++ b/tests/test_strings.py
@@ -2,13 +2,13 @@
 
 from __future__ import absolute_import
 
-import unittest
+import unittest2
 
 import rsa
 
 from tests.constants import unicode_string
 
-class StringTest(unittest.TestCase):
+class StringTest(unittest2.TestCase):
 
     def setUp(self):
         (self.pub, self.priv) = rsa.newkeys(384)
diff --git a/tests/test_varblock.py b/tests/test_varblock.py
index 6195258..24ea50f 100644
--- a/tests/test_varblock.py
+++ b/tests/test_varblock.py
@@ -1,20 +1,22 @@
 '''Tests varblock operations.'''
 
+
 try:
-    from StringIO import StringIO
+    from StringIO import StringIO as BytesIO
 except ImportError:
-    from io import StringIO
+    from io import BytesIO
 import unittest
 
 import rsa
+from rsa._compat import b
 from rsa import varblock
 
 class VarintTest(unittest.TestCase):
 
     def test_read_varint(self):
         
-        encoded = '\xac\x02crummy'
-        infile = StringIO(encoded)
+        encoded = b('\xac\x02crummy')
+        infile = BytesIO(encoded)
 
         (decoded, read) = varblock.read_varint(infile)
 
@@ -23,12 +25,12 @@ class VarintTest(unittest.TestCase):
         self.assertEqual(2, read)
 
         # The rest of the file should be untouched
-        self.assertEqual('crummy', infile.read())
+        self.assertEqual(b('crummy'), infile.read())
 
     def test_read_zero(self):
         
-        encoded = '\x00crummy'
-        infile = StringIO(encoded)
+        encoded = b('\x00crummy')
+        infile = BytesIO(encoded)
 
         (decoded, read) = varblock.read_varint(infile)
 
@@ -37,12 +39,12 @@ class VarintTest(unittest.TestCase):
         self.assertEqual(1, read)
 
         # The rest of the file should be untouched
-        self.assertEqual('crummy', infile.read())
+        self.assertEqual(b('crummy'), infile.read())
 
     def test_write_varint(self):
         
-        expected = '\xac\x02'
-        outfile = StringIO()
+        expected = b('\xac\x02')
+        outfile = BytesIO()
 
         written = varblock.write_varint(outfile, 300)
 
@@ -53,28 +55,28 @@ class VarintTest(unittest.TestCase):
 
     def test_write_zero(self):
         
-        outfile = StringIO()
+        outfile = BytesIO()
         written = varblock.write_varint(outfile, 0)
 
         # Test the returned values
-        self.assertEqual('\x00', outfile.getvalue())
+        self.assertEqual(b('\x00'), outfile.getvalue())
         self.assertEqual(1, written)
 
 
 class VarblockTest(unittest.TestCase):
 
     def test_yield_varblock(self):
-        infile = StringIO('\x01\x0512345\x06Sybren')
+        infile = BytesIO(b('\x01\x0512345\x06Sybren'))
 
         varblocks = list(varblock.yield_varblocks(infile))
-        self.assertEqual(['12345', 'Sybren'], varblocks)
+        self.assertEqual([b('12345'), b('Sybren')], varblocks)
 
 class FixedblockTest(unittest.TestCase):
 
     def test_yield_fixedblock(self):
 
-        infile = StringIO('123456Sybren')
+        infile = BytesIO(b('123456Sybren'))
 
         fixedblocks = list(varblock.yield_fixedblocks(infile, 6))
-        self.assertEqual(['123456', 'Sybren'], fixedblocks)
+        self.assertEqual([b('123456'), b('Sybren')], fixedblocks)