refactor test runner

2 files changed, 24 insertions, 12 deletions

diff --git a/psutil/tests/runner.py b/psutil/tests/runner.py
index 1c282f68..88bcd620 100755
--- a/psutil/tests/runner.py
+++ b/psutil/tests/runner.py
@@ -13,15 +13,25 @@ from psutil.tests import unittest
 from psutil.tests import VERBOSITY
 
 
-HERE = os.path.abspath(os.path.dirname(__file__))
-testmodules = [os.path.splitext(x)[0] for x in os.listdir(HERE)
-               if x.endswith('.py') and x.startswith('test_') and not
-               x.startswith('test_memory_leaks')]
-suite = unittest.TestSuite()
-for tm in testmodules:
-    # ...so that "make test" will print the full test paths
-    tm = "psutil.tests.%s" % tm
-    suite.addTest(unittest.defaultTestLoader.loadTestsFromName(tm))
-result = unittest.TextTestRunner(verbosity=VERBOSITY).run(suite)
-success = result.wasSuccessful()
-sys.exit(0 if success else 1)
+def get_suite():
+    HERE = os.path.abspath(os.path.dirname(__file__))
+    testmodules = [os.path.splitext(x)[0] for x in os.listdir(HERE)
+                   if x.endswith('.py') and x.startswith('test_') and not
+                   x.startswith('test_memory_leaks')]
+    suite = unittest.TestSuite()
+    for tm in testmodules:
+        # ...so that "make test" will print the full test paths
+        tm = "psutil.tests.%s" % tm
+        suite.addTest(unittest.defaultTestLoader.loadTestsFromName(tm))
+    return suite
+
+
+def main():
+    # run tests
+    result = unittest.TextTestRunner(verbosity=VERBOSITY).run(get_suite())
+    success = result.wasSuccessful()
+    sys.exit(0 if success else 1)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/setup.py b/setup.py
index 01543bee..47772da9 100755
--- a/setup.py
+++ b/setup.py
@@ -270,6 +270,8 @@ def main():
         license='BSD',
         packages=['psutil', 'psutil.tests'],
         ext_modules=extensions,
+        test_suite="psutil.tests.runner.get_suite",
+        tests_require=['ipaddress', 'mock', 'unittest2'],
         # see: python setup.py register --list-classifiers
         classifiers=[
             'Development Status :: 5 - Production/Stable',