diff --git a/setup.py b/setup.py
index 1f633d414d933652649723823efea177c22e68d5..413994a28281d19fec0f71ad7bd68e0c57cd402f 100755
--- a/setup.py
+++ b/setup.py
@@ -25,10 +25,7 @@ from distutils import log
from distutils.command.install import install
from distutils.core import setup
from distutils.core import Command
-if sys.version_info < (2, 7):
- import unittest2 as unittest
-else:
- import unittest
+
BASE_DIR = os.path.dirname(globals().get('__file__', os.getcwd()))
@@ -44,6 +41,14 @@ class TestCmd(Command):
pass
def run(self):
+ if sys.version_info < (2, 7):
+ try:
+ import unittest2 as unittest
+ except ImportError:
+ raise RuntimeError('unittest2 required for running tests under Python < 2.7.')
+ else:
+ import unittest
+
test_dir = os.path.join(BASE_DIR, 'tests')
tests = unittest.TestLoader().discover(test_dir)
@@ -52,6 +57,7 @@ class TestCmd(Command):
if not result.wasSuccessful():
sys.exit(1)
+
class CleanCmd(Command):
description = 'Remove all generated files.'
user_options = []
@@ -153,6 +159,7 @@ class InstallCmd(install):
os.path.join(manpage_dir, manpage_file),
dry_run=self.dry_run)
+
# Only override install if not being run by setuptools.
cmdclass = {'test': TestCmd,
'dist_clean': CleanCmd,
@@ -160,6 +167,7 @@ cmdclass = {'test': TestCmd,
if 'setuptools' not in dir():
cmdclass['install'] = InstallCmd
+
setup(
cmdclass=cmdclass,
name='chkcrontab',