--- a/test/test_pytest/test_flaky_pytest_plugin.py
+++ b/test/test_pytest/test_flaky_pytest_plugin.py
@@ -2,7 +2,10 @@
 
 from __future__ import unicode_literals
 from io import StringIO
-from mock import Mock, patch
+try:
+    from unittest.mock import Mock, patch
+except ImportError:
+    from mock import Mock, patch
 # pylint:disable=import-error
 import pytest
 from _pytest.runner import CallInfo