Use base unittest or unittest2 depending on python version

1 files changed, 2 insertions, 2 deletions

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' ])