Index: kafka-python-2.0.2/test/test_consumer_integration.py
===================================================================
--- kafka-python-2.0.2.orig/test/test_consumer_integration.py
+++ kafka-python-2.0.2/test/test_consumer_integration.py
@@ -1,7 +1,10 @@
 import logging
 import time
 
-from mock import patch
+try:
+    from unittest.mock import patch
+except ImportError:
+    from mock import patch
 import pytest
 from kafka.vendor.six.moves import range
 
Index: kafka-python-2.0.2/test/test_conn.py
===================================================================
--- kafka-python-2.0.2.orig/test/test_conn.py
+++ kafka-python-2.0.2/test/test_conn.py
@@ -4,7 +4,10 @@ from __future__ import absolute_import
 from errno import EALREADY, EINPROGRESS, EISCONN, ECONNRESET
 import socket
 
-import mock
+try:
+    from unittest import mock
+except ImportError:
+    import mock
 import pytest
 
 from kafka.conn import BrokerConnection, ConnectionStates, collect_hosts