diff options
| -rw-r--r-- | tests/compat.py | 6 | ||||
| -rw-r--r-- | tests/test_declarations.py | 6 |
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') |