diff --git a/test/test_context.py b/test/test_context.py
index 8c3fd856d..da9b22f65 100644
--- a/test/test_context.py
+++ b/test/test_context.py
@@ -1,7 +1,7 @@
 """
 OffsetCommitContext tests.
 """
-from unittest2 import TestCase
+from . import unittest
 
 from mock import MagicMock, patch
 
@@ -9,7 +9,7 @@
 from kafka.context import OffsetCommitContext
 
 
-class TestOffsetCommitContext(TestCase):
+class TestOffsetCommitContext(unittest.TestCase):
     """
     OffsetCommitContext tests.
     """