2 files changed, 8 insertions, 4 deletions
diff --git a/tests/compat.py b/tests/compat.py
index 8d4f1d0..769ffd4 100644
--- a/tests/compat.py
+++ b/tests/compat.py
@@ -29,3 +29,9 @@ try:
     import unittest2 as unittest
 except ImportError:
     import unittest
+
+if is_python2:
+    import mock
+else:
+    from unittest import mock
+
diff --git a/tests/test_declarations.py b/tests/test_declarations.py
index 59a3955..93e11d0 100644
--- a/tests/test_declarations.py
+++ b/tests/test_declarations.py
@@ -24,11 +24,9 @@ import datetime
 import itertools
 import warnings
 
-from mock import MagicMock
-
 from factory import declarations
 
-from .compat import unittest
+from .compat import mock, unittest
 from . import tools
 
 
@@ -299,7 +297,7 @@ class RelatedFactoryTestCase(unittest.TestCase):
 
 class PostGenerationMethodCallTestCase(unittest.TestCase):
     def setUp(self):
-        self.obj = MagicMock()
+        self.obj = mock.MagicMock()
 
     def test_simplest_setup_and_call(self):
         decl = declarations.PostGenerationMethodCall('method')