unify tests for py2 and py3

1 files changed, 6 insertions, 9 deletions

diff --git a/test/test_obj.py b/test/test_obj.py
index 0eebe7b..6357cfc 100644
--- a/test/test_obj.py
+++ b/test/test_obj.py
@@ -7,22 +7,22 @@ from nose.tools import *
 from msgpack import packs, unpacks
 
 def _decode_complex(obj):
-    if '__complex__' in obj:
-        return complex(obj['real'], obj['imag'])
+    if b'__complex__' in obj:
+        return complex(obj[b'real'], obj[b'imag'])
     return obj
 
 def _encode_complex(obj):
     if isinstance(obj, complex):
-        return {'__complex__': True, 'real': 1, 'imag': 2}
+        return {b'__complex__': True, b'real': 1, b'imag': 2}
     return obj
 
 def test_encode_hook():
     packed = packs([3, 1+2j], default=_encode_complex)
     unpacked = unpacks(packed)
-    eq_(unpacked[1], {'__complex__': True, 'real': 1, 'imag': 2})
+    eq_(unpacked[1], {b'__complex__': True, b'real': 1, b'imag': 2})
 
 def test_decode_hook():
-    packed = packs([3, {'__complex__': True, 'real': 1, 'imag': 2}])
+    packed = packs([3, {b'__complex__': True, b'real': 1, b'imag': 2}])
     unpacked = unpacks(packed, object_hook=_decode_complex)
     eq_(unpacked[1], 1+2j)
 
@@ -40,7 +40,4 @@ def test_array_hook():
     eq_(unpacked, '123')
 
 if __name__ == '__main__':
-    test_decode_hook()
-    test_encode_hook()
-    test_bad_hook()
-    test_array_hook()
+    main()