import unittest.mock as mock

2 files changed, 10 insertions, 2 deletions

diff --git a/tests/integration/tests.py b/tests/integration/tests.py
index 79fadda..024feab 100644
--- a/tests/integration/tests.py
+++ b/tests/integration/tests.py
@@ -99,7 +99,11 @@ from pokemontools.crystal import (
 import pokemontools.wram
 
 import unittest
-import mock
+
+try:
+    import unittest.mock as mock
+except ImportError:
+    import mock
 
 class BasicTestCase(unittest.TestCase):
     "this is where i cram all of my unit tests together"
diff --git a/tests/tests.py b/tests/tests.py
index 4c7f7a4..13b1e49 100644
--- a/tests/tests.py
+++ b/tests/tests.py
@@ -95,7 +95,11 @@ from pokemontools.crystal import (
 )
 
 import unittest
-import mock
+
+try:
+    import unittest.mock as mock
+except ImportError:
+    import mock
 
 class TestCram(unittest.TestCase):
     "this is where i cram all of my unit tests together"