diff --git a/tests/test_progressbar.py b/tests/test_progressbar.py
index 304e14e..13b2427 100644
--- a/tests/test_progressbar.py
+++ b/tests/test_progressbar.py
@@ -28,10 +28,14 @@
 Test the progressbar 'FileTransferSpeed' integration
 """
 
+try:
+    import unittest.mock
+except ImportError:
+    import mock
+
 from . import TestCase
 import bitmath
 from bitmath.integrations.bmprogressbar import BitmathFileTransferSpeed
-import mock
 import progressbar
 
 
diff --git a/tests/test_query_device_capacity.py b/tests/test_query_device_capacity.py
index 64ca99e..4ddda41 100644
--- a/tests/test_query_device_capacity.py
+++ b/tests/test_query_device_capacity.py
@@ -30,9 +30,13 @@
 
 from . import TestCase
 import bitmath
-import mock
 import struct
 
+try:
+    import unittest.mock
+except ImportError:
+    import mock
+
 try:
     # Python 3.3+
     from contextlib import ExitStack, contextmanager