diff --git a/fire/core_test.py b/fire/core_test.py index 27c9f418..71e67e9a 100644 --- a/fire/core_test.py +++ b/fire/core_test.py @@ -22,7 +22,11 @@ from fire import test_components as tc from fire import testutils from fire import trace -import mock + +try: + import mock # python 2.x +except ModuleNotFoundError: + from unittest import mock # python 3.x import six diff --git a/fire/fire_import_test.py b/fire/fire_import_test.py index c5975681..3f379eff 100644 --- a/fire/fire_import_test.py +++ b/fire/fire_import_test.py @@ -18,7 +18,10 @@ import fire from fire import testutils -import mock +try: + import mock # python 2.x +except ModuleNotFoundError: + from unittest import mock # python 3.x class FireImportTest(testutils.BaseTestCase): diff --git a/fire/fire_test.py b/fire/fire_test.py index 4f2ffb9b..69b58c76 100644 --- a/fire/fire_test.py +++ b/fire/fire_test.py @@ -25,7 +25,10 @@ from fire import test_components as tc from fire import testutils -import mock +try: + import mock # python 2.x +except ModuleNotFoundError: + from unittest import mock # python 3.x import six diff --git a/fire/interact_test.py b/fire/interact_test.py index 29fa7597..7dea1f0e 100644 --- a/fire/interact_test.py +++ b/fire/interact_test.py @@ -21,7 +21,10 @@ from fire import interact from fire import testutils -import mock +try: + import mock # python 2.x +except ModuleNotFoundError: + from unittest import mock # python 3.x try: diff --git a/fire/testutils.py b/fire/testutils.py index 3463fd50..628617d6 100644 --- a/fire/testutils.py +++ b/fire/testutils.py @@ -26,7 +26,10 @@ from fire import core from fire import trace -import mock +try: + import mock # python 2.x +except ModuleNotFoundError: + from unittest import mock # python 3.x import six