need to use SQLA's approach of calling mock from unittest for py3k

2 files changed, 13 insertions, 2 deletions

diff --git a/tests/__init__.py b/tests/__init__.py
index 2fb34da..d6f686c 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -14,7 +14,7 @@ from sqlalchemy.util import decorator
 
 import alembic
 from alembic import util
-from alembic.compat import string_types, text_type, u
+from alembic.compat import string_types, text_type, u, py33
 from alembic.migration import MigrationContext
 from alembic.environment import EnvironmentContext
 from alembic.operations import Operations
@@ -27,6 +27,17 @@ files_directory = os.path.join(os.path.dirname(__file__), 'files')
 testing_config = configparser.ConfigParser()
 testing_config.read(['test.cfg'])
 
+if py33:
+    from unittest.mock import Mock, call
+else:
+    try:
+        from mock import Mock, call
+    except ImportError:
+        raise ImportError(
+                "Alembic's test suite requires the "
+                "'mock' library as of 0.6.1.")
+
+
 def sqlite_db():
     # sqlite caches table pragma info
     # per connection, so create a new
diff --git a/tests/test_config.py b/tests/test_config.py
index 3f7862c..6164eb9 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -5,7 +5,7 @@ from alembic.migration import MigrationContext
 from alembic.operations import Operations
 from alembic.script import ScriptDirectory
 import unittest
-from mock import Mock, call
+from . import Mock, call
 
 from . import eq_, capture_db, assert_raises_message