Use base unittest or unittest2 depending on python version

1 files changed, 6 insertions, 1 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"],