Use base unittest or unittest2 depending on python version

13 files changed, 43 insertions, 30 deletions

diff --git a/setup.py b/setup.py
index a7e1400..e16fe0e 100644
--- a/setup.py
+++ b/setup.py
@@ -22,11 +22,16 @@ class Tox(Command):
         sys.exit(tox.cmdline([]))
 
 
+test_require = ['tox', 'mock']
+if sys.version_info < (2, 7):
+    test_require.append('unittest2')
+
+
 setup(
     name="kafka-python",
     version=__version__,
 
-    tests_require=["tox", "mock", "unittest2"],
+    tests_require=test_require,
     cmdclass={"test": Tox},
 
     packages=["kafka"],
diff --git a/test/__init__.py b/test/__init__.py
index e69de29..c4d1e80 100644
--- a/test/__init__.py
+++ b/test/__init__.py
@@ -0,0 +1,6 @@
+import sys
+
+if sys.version_info < (2, 7):
+    import unittest2 as unittest
+else:
+    import unittest
diff --git a/test/test_client.py b/test/test_client.py
index 32a2256..b7591e9 100644
--- a/test/test_client.py
+++ b/test/test_client.py
@@ -1,4 +1,4 @@
-import unittest2
+from . import unittest
 
 from mock import MagicMock, patch
 
@@ -10,7 +10,7 @@ from kafka.common import (
 )
 from kafka.protocol import create_message
 
-class TestKafkaClient(unittest2.TestCase):
+class TestKafkaClient(unittest.TestCase):
     def test_init_with_list(self):
         with patch.object(KafkaClient, 'load_metadata_for_topics'):
             client = KafkaClient(hosts=['kafka01:9092', 'kafka02:9092', 'kafka03:9092'])
diff --git a/test/test_client_integration.py b/test/test_client_integration.py
index b5bcb22..3eb917f 100644
--- a/test/test_client_integration.py
+++ b/test/test_client_integration.py
@@ -1,6 +1,6 @@
 import os
 import socket
-import unittest2
+from . import unittest
 
 import kafka
 from kafka.common import *
@@ -24,7 +24,7 @@ class TestKafkaClientIntegration(KafkaIntegrationTestCase):
         cls.server.close()
         cls.zk.close()
 
-    @unittest2.skip("This doesn't appear to work on Linux?")
+    @unittest.skip("This doesn't appear to work on Linux?")
     def test_timeout(self):
         server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         server_port = get_open_port()
diff --git a/test/test_codec.py b/test/test_codec.py
index 2e6f67e..c409052 100644
--- a/test/test_codec.py
+++ b/test/test_codec.py
@@ -1,5 +1,5 @@
 import struct
-import unittest2
+from . import unittest
 
 from kafka.codec import (
     has_snappy, gzip_encode, gzip_decode,
@@ -10,21 +10,21 @@ from kafka.protocol import (
 )
 from testutil import *
 
-class TestCodec(unittest2.TestCase):
+class TestCodec(unittest.TestCase):
     def test_gzip(self):
         for i in xrange(1000):
             s1 = random_string(100)
             s2 = gzip_decode(gzip_encode(s1))
             self.assertEquals(s1, s2)
 
-    @unittest2.skipUnless(has_snappy(), "Snappy not available")
+    @unittest.skipUnless(has_snappy(), "Snappy not available")
     def test_snappy(self):
         for i in xrange(1000):
             s1 = random_string(100)
             s2 = snappy_decode(snappy_encode(s1))
             self.assertEquals(s1, s2)
 
-    @unittest2.skipUnless(has_snappy(), "Snappy not available")
+    @unittest.skipUnless(has_snappy(), "Snappy not available")
     def test_snappy_detect_xerial(self):
         import kafka as kafka1
         _detect_xerial_stream = kafka1.codec._detect_xerial_stream
@@ -41,7 +41,7 @@ class TestCodec(unittest2.TestCase):
         self.assertFalse(_detect_xerial_stream(random_snappy))
         self.assertFalse(_detect_xerial_stream(short_data))
 
-    @unittest2.skipUnless(has_snappy(), "Snappy not available")
+    @unittest.skipUnless(has_snappy(), "Snappy not available")
     def test_snappy_decode_xerial(self):
         header = b'\x82SNAPPY\x00\x00\x00\x00\x01\x00\x00\x00\x01'
         random_snappy = snappy_encode('SNAPPY' * 50)
@@ -55,7 +55,7 @@ class TestCodec(unittest2.TestCase):
 
         self.assertEquals(snappy_decode(to_test), ('SNAPPY' * 50) + ('XERIAL' * 50))
 
-    @unittest2.skipUnless(has_snappy(), "Snappy not available")
+    @unittest.skipUnless(has_snappy(), "Snappy not available")
     def test_snappy_encode_xerial(self):
         to_ensure = b'\x82SNAPPY\x00\x00\x00\x00\x01\x00\x00\x00\x01' + \
             '\x00\x00\x00\x18' + \
diff --git a/test/test_conn.py b/test/test_conn.py
index 5451398..bc644d7 100644
--- a/test/test_conn.py
+++ b/test/test_conn.py
@@ -2,12 +2,12 @@ import socket
 import struct
 
 import mock
-import unittest2
+from . import unittest
 
 from kafka.common import *
 from kafka.conn import *
 
-class ConnTest(unittest2.TestCase):
+class ConnTest(unittest.TestCase):
     def setUp(self):
         self.config = {
             'host': 'localhost',
diff --git a/test/test_consumer.py b/test/test_consumer.py
index 778d76a..7d7fc46 100644
--- a/test/test_consumer.py
+++ b/test/test_consumer.py
@@ -1,7 +1,7 @@
 import os
 import random
 import struct
-import unittest2
+from . import unittest
 
 from mock import MagicMock, patch
 
@@ -16,7 +16,7 @@ from kafka.protocol import (
     create_message, KafkaProtocol
 )
 
-class TestKafkaConsumer(unittest2.TestCase):
+class TestKafkaConsumer(unittest.TestCase):
     def test_non_integer_partitions(self):
         with self.assertRaises(AssertionError):
             consumer = SimpleConsumer(MagicMock(), 'group', 'topic', partitions = [ '0' ])
diff --git a/test/test_failover_integration.py b/test/test_failover_integration.py
index 6c0e662..40269ee 100644
--- a/test/test_failover_integration.py
+++ b/test/test_failover_integration.py
@@ -1,7 +1,7 @@
 import logging
 import os
 import time
-import unittest2
+from . import unittest
 
 from kafka import *  # noqa
 from kafka.common import *  # noqa
@@ -82,7 +82,7 @@ class TestFailover(KafkaIntegrationTestCase):
 
 
     #@kafka_versions("all")
-    @unittest2.skip("async producer does not support reliable failover yet")
+    @unittest.skip("async producer does not support reliable failover yet")
     def test_switch_leader_async(self):
         key, topic, partition = random_string(5), self.topic, 0
 
diff --git a/test/test_package.py b/test/test_package.py
index a6a3a14..9b69a7c 100644
--- a/test/test_package.py
+++ b/test/test_package.py
@@ -1,6 +1,6 @@
-import unittest2
+from . import unittest
 
-class TestPackage(unittest2.TestCase):
+class TestPackage(unittest.TestCase):
     def test_top_level_namespace(self):
         import kafka as kafka1
         self.assertEquals(kafka1.KafkaClient.__name__, "KafkaClient")
diff --git a/test/test_producer.py b/test/test_producer.py
index a84e20f..9d21369 100644
--- a/test/test_producer.py
+++ b/test/test_producer.py
@@ -4,14 +4,14 @@ import logging
 import os
 import random
 import struct
-import unittest2
+from . import unittest
 
 from mock import MagicMock, patch
 
 from kafka import KafkaClient
 from kafka.producer import Producer
 
-class TestKafkaProducer(unittest2.TestCase):
+class TestKafkaProducer(unittest.TestCase):
     def test_producer_message_types(self):
 
         producer = Producer(MagicMock())
diff --git a/test/test_protocol.py b/test/test_protocol.py
index 2089f48..e586228 100644
--- a/test/test_protocol.py
+++ b/test/test_protocol.py
@@ -1,7 +1,7 @@
 import contextlib
 from contextlib import contextmanager
 import struct
-import unittest2
+from . import unittest
 
 import mock
 from mock import sentinel
@@ -27,7 +27,7 @@ from kafka.protocol import (
     create_message_set
 )
 
-class TestProtocol(unittest2.TestCase):
+class TestProtocol(unittest.TestCase):
     def test_create_message(self):
         payload = "test"
         key = "key"
@@ -65,7 +65,7 @@ class TestProtocol(unittest2.TestCase):
 
         self.assertEqual(decoded, expect)
 
-    @unittest2.skipUnless(has_snappy(), "Snappy not available")
+    @unittest.skipUnless(has_snappy(), "Snappy not available")
     def test_create_snappy(self):
         payloads = ["v1", "v2"]
         msg = create_snappy_message(payloads)
@@ -222,7 +222,7 @@ class TestProtocol(unittest2.TestCase):
         self.assertEqual(returned_offset2, 0)
         self.assertEqual(decoded_message2, create_message("v2"))
 
-    @unittest2.skipUnless(has_snappy(), "Snappy not available")
+    @unittest.skipUnless(has_snappy(), "Snappy not available")
     def test_decode_message_snappy(self):
         snappy_encoded = ('\xec\x80\xa1\x95\x00\x02\xff\xff\xff\xff\x00\x00'
                           '\x00,8\x00\x00\x19\x01@\x10L\x9f[\xc2\x00\x00\xff'
diff --git a/test/test_util.py b/test/test_util.py
index 7b5f294..4772d3a 100644
--- a/test/test_util.py
+++ b/test/test_util.py
@@ -1,12 +1,13 @@
 # -*- coding: utf-8 -*-
 import struct
-import unittest2
 import kafka.util
 import kafka.common
 
+from . import unittest
 
-class UtilTest(unittest2.TestCase):
-    @unittest2.skip("Unwritten")
+
+class UtilTest(unittest.TestCase):
+    @unittest.skip("Unwritten")
     def test_relative_unpack(self):
         pass
 
diff --git a/test/testutil.py b/test/testutil.py
index dc8eea0..e7dcda6 100644
--- a/test/testutil.py
+++ b/test/testutil.py
@@ -5,9 +5,10 @@ import random
 import socket
 import string
 import time
-import unittest2
 import uuid
 
+from . import unittest
+
 from kafka.common import OffsetRequest
 from kafka import KafkaClient
 
@@ -45,7 +46,7 @@ def get_open_port():
     sock.close()
     return port
 
-class KafkaIntegrationTestCase(unittest2.TestCase):
+class KafkaIntegrationTestCase(unittest.TestCase):
     create_client = True
     topic = None
     server = None