From c4467c90ef53cd9eba8ab8be620d71b8477d6ce9 Mon Sep 17 00:00:00 2001
From: Giampaolo Rodola <g.rodola@gmail.com>
Date: Wed, 13 Mar 2019 07:30:52 +0100
Subject: [PATCH] Coloured tests (#1459)

Fixes #1458.
---
 MANIFEST.in                       |   1 +
 Makefile                          |   4 +-
 psutil/tests/__init__.py          |  50 +-----------
 psutil/tests/__main__.py          |   2 +-
 psutil/tests/runner.py            | 122 ++++++++++++++++++++++++++++++
 psutil/tests/test_aix.py          |   2 +-
 psutil/tests/test_bsd.py          |   2 +-
 psutil/tests/test_connections.py  |   2 +-
 psutil/tests/test_contracts.py    |   2 +-
 psutil/tests/test_linux.py        |   2 +-
 psutil/tests/test_memory_leaks.py |   2 +-
 psutil/tests/test_misc.py         |   2 +-
 psutil/tests/test_osx.py          |   2 +-
 psutil/tests/test_posix.py        |   2 +-
 psutil/tests/test_process.py      |   2 +-
 psutil/tests/test_sunos.py        |   2 +-
 psutil/tests/test_system.py       |   2 +-
 psutil/tests/test_unicode.py      |   2 +-
 psutil/tests/test_windows.py      |   2 +-
 19 files changed, 144 insertions(+), 63 deletions(-)
 create mode 100755 psutil/tests/runner.py

diff --git a/MANIFEST.in b/MANIFEST.in
index 3e7b5e7d5..6ab7c6483 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -80,6 +80,7 @@ include psutil/arch/windows/services.h
 include psutil/tests/README.rst
 include psutil/tests/__init__.py
 include psutil/tests/__main__.py
+include psutil/tests/runner.py
 include psutil/tests/test_aix.py
 include psutil/tests/test_bsd.py
 include psutil/tests/test_connections.py
diff --git a/Makefile b/Makefile
index 59ac15068..7d838e876 100644
--- a/Makefile
+++ b/Makefile
@@ -113,11 +113,11 @@ test:  ## Run all tests.
 
 test-process:  ## Run process-related API tests.
 	${MAKE} install
-	$(TEST_PREFIX) $(PYTHON) -m unittest -v psutil.tests.test_process
+	$(TEST_PREFIX) $(PYTHON) psutil/tests/test_process.py
 
 test-system:  ## Run system-related API tests.
 	${MAKE} install
-	$(TEST_PREFIX) $(PYTHON) -m unittest -v psutil.tests.test_system
+	$(TEST_PREFIX) $(PYTHON) psutil/tests/test_system.py
 
 test-misc:  ## Run miscellaneous tests.
 	${MAKE} install
diff --git a/psutil/tests/__init__.py b/psutil/tests/__init__.py
index 74b0c969a..0a29c59d9 100644
--- a/psutil/tests/__init__.py
+++ b/psutil/tests/__init__.py
@@ -69,7 +69,6 @@
     'APPVEYOR', 'DEVNULL', 'GLOBAL_TIMEOUT', 'MEMORY_TOLERANCE', 'NO_RETRIES',
     'PYPY', 'PYTHON_EXE', 'ROOT_DIR', 'SCRIPTS_DIR', 'TESTFILE_PREFIX',
     'TESTFN', 'TESTFN_UNICODE', 'TOX', 'TRAVIS', 'VALID_PROC_STATUSES',
-    'VERBOSITY',
     "HAS_CPU_AFFINITY", "HAS_CPU_FREQ", "HAS_ENVIRON", "HAS_PROC_IO_COUNTERS",
     "HAS_IONICE", "HAS_MEMORY_MAPS", "HAS_PROC_CPU_NUM", "HAS_RLIMIT",
     "HAS_SENSORS_BATTERY", "HAS_BATTERY", "HAS_SENSORS_FANS",
@@ -127,8 +126,6 @@
 MEMORY_TOLERANCE = 500 * 1024  # 500KB
 # the timeout used in functions which have to wait
 GLOBAL_TIMEOUT = 3 if TRAVIS or APPVEYOR else 0.5
-# test output verbosity
-VERBOSITY = 1 if os.getenv('SILENT') or TOX else 2
 # be more tolerant if we're on travis / appveyor in order to avoid
 # false positives
 if TRAVIS or APPVEYOR:
@@ -802,9 +799,11 @@ class TestCase(unittest.TestCase):
     # Print a full path representation of the single unit tests
     # being run.
     def __str__(self):
+        fqmod = self.__class__.__module__
+        if not fqmod.startswith('psutil.'):
+            fqmod = 'psutil.tests.' + fqmod
         return "%s.%s.%s" % (
-            self.__class__.__module__, self.__class__.__name__,
-            self._testMethodName)
+            fqmod, self.__class__.__name__, self._testMethodName)
 
     # assertRaisesRegexp renamed to assertRaisesRegex in 3.3;
     # add support for the new name.
@@ -816,47 +815,6 @@ def __str__(self):
 unittest.TestCase = TestCase
 
 
-def _setup_tests():
-    if 'PSUTIL_TESTING' not in os.environ:
-        # This won't work on Windows but set_testing() below will do it.
-        os.environ['PSUTIL_TESTING'] = '1'
-    psutil._psplatform.cext.set_testing()
-
-
-def get_suite():
-    testmods = [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')]
-    if "WHEELHOUSE_UPLOADER_USERNAME" in os.environ:
-        testmods = [x for x in testmods if not x.endswith((
-                    "osx", "posix", "linux"))]
-    suite = unittest.TestSuite()
-    for tm in testmods:
-        # ...so that the full test paths are printed on screen
-        tm = "psutil.tests.%s" % tm
-        suite.addTest(unittest.defaultTestLoader.loadTestsFromName(tm))
-    return suite
-
-
-def run_suite():
-    _setup_tests()
-    result = unittest.TextTestRunner(verbosity=VERBOSITY).run(get_suite())
-    success = result.wasSuccessful()
-    sys.exit(0 if success else 1)
-
-
-def run_test_module_by_name(name):
-    # testmodules = [os.path.splitext(x)[0] for x in os.listdir(HERE)
-    #                if x.endswith('.py') and x.startswith('test_')]
-    _setup_tests()
-    name = os.path.splitext(os.path.basename(name))[0]
-    suite = unittest.TestSuite()
-    suite.addTest(unittest.defaultTestLoader.loadTestsFromName(name))
-    result = unittest.TextTestRunner(verbosity=VERBOSITY).run(suite)
-    success = result.wasSuccessful()
-    sys.exit(0 if success else 1)
-
-
 def retry_before_failing(retries=NO_RETRIES):
     """Decorator which runs a test function and retries N times before
     actually failing.
diff --git a/psutil/tests/__main__.py b/psutil/tests/__main__.py
index 36554a126..3180c279c 100755
--- a/psutil/tests/__main__.py
+++ b/psutil/tests/__main__.py
@@ -21,7 +21,7 @@
     from urllib2 import urlopen
 
 from psutil.tests import PYTHON_EXE
-from psutil.tests import run_suite
+from psutil.tests.runner import run_suite
 
 
 HERE = os.path.abspath(os.path.dirname(__file__))
diff --git a/psutil/tests/runner.py b/psutil/tests/runner.py
new file mode 100755
index 000000000..33f960002
--- /dev/null
+++ b/psutil/tests/runner.py
@@ -0,0 +1,122 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""
+Unit test runner, providing colourized output.
+"""
+
+import os
+import sys
+import unittest
+from unittest import TestResult
+from unittest import TextTestResult
+from unittest import TextTestRunner
+
+import psutil
+from psutil.tests import TOX
+
+
+HERE = os.path.abspath(os.path.dirname(__file__))
+VERBOSITY = 1 if TOX else 2
+GREEN = 1
+RED = 2
+BROWN = 94
+
+
+def term_supports_colors(file=sys.stdout):
+    try:
+        import curses
+        assert file.isatty()
+        curses.setupterm()
+        assert curses.tigetnum("colors") > 0
+    except Exception:
+        return False
+    else:
+        return True
+
+
+def hilite(s, color, bold=False):
+    """Return an highlighted version of 'string'."""
+    attr = []
+    if color == GREEN:
+        attr.append('32')
+    elif color == RED:
+        attr.append('91')
+    elif color == BROWN:
+        attr.append('33')
+    else:
+        raise ValueError("unrecognized color")
+    if bold:
+        attr.append('1')
+    return '\x1b[%sm%s\x1b[0m' % (';'.join(attr), s)
+
+
+class ColouredResult(unittest.TextTestResult):
+
+    def addSuccess(self, test):
+        TestResult.addSuccess(self, test)
+        self.stream.writeln(hilite("OK", GREEN))
+
+    def addError(self, test, err):
+        TestResult.addError(self, test, err)
+        self.stream.writeln(hilite("ERROR", RED, bold=True))
+
+    def addFailure(self, test, err):
+        TestResult.addFailure(self, test, err)
+        self.stream.writeln(hilite("FAIL", RED))
+
+    def addSkip(self, test, reason):
+        TestResult.addSkip(self, test, reason)
+        self.stream.writeln(hilite("skipped: %s" % reason, BROWN))
+
+    def printErrorList(self, flavour, errors):
+        flavour = hilite(flavour, RED)
+        TextTestResult.printErrorList(self, flavour, errors)
+
+
+class ColouredRunner(TextTestRunner):
+    resultclass = ColouredResult if term_supports_colors() else TextTestResult
+
+
+def setup_tests():
+    if 'PSUTIL_TESTING' not in os.environ:
+        # This won't work on Windows but set_testing() below will do it.
+        os.environ['PSUTIL_TESTING'] = '1'
+    psutil._psplatform.cext.set_testing()
+
+
+def get_suite():
+    testmods = [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')]
+    if "WHEELHOUSE_UPLOADER_USERNAME" in os.environ:
+        testmods = [x for x in testmods if not x.endswith((
+                    "osx", "posix", "linux"))]
+    suite = unittest.TestSuite()
+    for tm in testmods:
+        # ...so that the full test paths are printed on screen
+        tm = "psutil.tests.%s" % tm
+        suite.addTest(unittest.defaultTestLoader.loadTestsFromName(tm))
+    return suite
+
+
+def run_test_module_by_name(name):
+    # testmodules = [os.path.splitext(x)[0] for x in os.listdir(HERE)
+    #                if x.endswith('.py') and x.startswith('test_')]
+    setup_tests()
+    name = os.path.splitext(os.path.basename(name))[0]
+    suite = unittest.TestSuite()
+    suite.addTest(unittest.defaultTestLoader.loadTestsFromName(name))
+    result = ColouredRunner(verbosity=VERBOSITY).run(suite)
+    success = result.wasSuccessful()
+    sys.exit(0 if success else 1)
+
+
+def run_suite():
+    setup_tests()
+    result = ColouredRunner(verbosity=VERBOSITY).run(get_suite())
+    success = result.wasSuccessful()
+    sys.exit(0 if success else 1)
diff --git a/psutil/tests/test_aix.py b/psutil/tests/test_aix.py
index 4f0da4e00..5294493b2 100755
--- a/psutil/tests/test_aix.py
+++ b/psutil/tests/test_aix.py
@@ -11,7 +11,6 @@
 import re
 
 from psutil import AIX
-from psutil.tests import run_test_module_by_name
 from psutil.tests import sh
 from psutil.tests import unittest
 import psutil
@@ -118,4 +117,5 @@ def test_net_if_addrs_names(self):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)
diff --git a/psutil/tests/test_bsd.py b/psutil/tests/test_bsd.py
index 34b66ca6b..d06bfef91 100755
--- a/psutil/tests/test_bsd.py
+++ b/psutil/tests/test_bsd.py
@@ -25,7 +25,6 @@
 from psutil.tests import MEMORY_TOLERANCE
 from psutil.tests import reap_children
 from psutil.tests import retry_before_failing
-from psutil.tests import run_test_module_by_name
 from psutil.tests import sh
 from psutil.tests import unittest
 from psutil.tests import which
@@ -550,4 +549,5 @@ def test_cpu_stats_ctx_switches(self):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)
diff --git a/psutil/tests/test_connections.py b/psutil/tests/test_connections.py
index c97fc409e..2a79338ea 100755
--- a/psutil/tests/test_connections.py
+++ b/psutil/tests/test_connections.py
@@ -35,7 +35,6 @@
 from psutil.tests import HAS_CONNECTIONS_UNIX
 from psutil.tests import pyrun
 from psutil.tests import reap_children
-from psutil.tests import run_test_module_by_name
 from psutil.tests import safe_rmpath
 from psutil.tests import skip_on_access_denied
 from psutil.tests import tcp_socketpair
@@ -523,4 +522,5 @@ def test_connection_constants(self):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)
diff --git a/psutil/tests/test_contracts.py b/psutil/tests/test_contracts.py
index 90fa8a155..892ffb9b3 100755
--- a/psutil/tests/test_contracts.py
+++ b/psutil/tests/test_contracts.py
@@ -38,7 +38,6 @@
 from psutil.tests import HAS_SENSORS_FANS
 from psutil.tests import HAS_SENSORS_TEMPERATURES
 from psutil.tests import is_namedtuple
-from psutil.tests import run_test_module_by_name
 from psutil.tests import safe_rmpath
 from psutil.tests import skip_on_access_denied
 from psutil.tests import TESTFN
@@ -654,4 +653,5 @@ def environ(self, ret, proc):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)
diff --git a/psutil/tests/test_linux.py b/psutil/tests/test_linux.py
index 0f981b6bc..de96f049c 100755
--- a/psutil/tests/test_linux.py
+++ b/psutil/tests/test_linux.py
@@ -38,7 +38,6 @@
 from psutil.tests import reap_children
 from psutil.tests import reload_module
 from psutil.tests import retry_before_failing
-from psutil.tests import run_test_module_by_name
 from psutil.tests import safe_rmpath
 from psutil.tests import sh
 from psutil.tests import skip_on_not_implemented
@@ -2080,4 +2079,5 @@ def test_cat(self):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)
diff --git a/psutil/tests/test_memory_leaks.py b/psutil/tests/test_memory_leaks.py
index cbe37d3e6..f90e8377c 100755
--- a/psutil/tests/test_memory_leaks.py
+++ b/psutil/tests/test_memory_leaks.py
@@ -46,7 +46,6 @@
 from psutil.tests import HAS_SENSORS_FANS
 from psutil.tests import HAS_SENSORS_TEMPERATURES
 from psutil.tests import reap_children
-from psutil.tests import run_test_module_by_name
 from psutil.tests import safe_rmpath
 from psutil.tests import skip_on_access_denied
 from psutil.tests import TESTFN
@@ -607,4 +606,5 @@ def test_win_service_get_description(self):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)
diff --git a/psutil/tests/test_misc.py b/psutil/tests/test_misc.py
index 04b45948e..36f8bf06f 100755
--- a/psutil/tests/test_misc.py
+++ b/psutil/tests/test_misc.py
@@ -55,7 +55,6 @@
 from psutil.tests import reload_module
 from psutil.tests import retry
 from psutil.tests import ROOT_DIR
-from psutil.tests import run_test_module_by_name
 from psutil.tests import safe_mkdir
 from psutil.tests import safe_rmpath
 from psutil.tests import SCRIPTS_DIR
@@ -1054,4 +1053,5 @@ def test_is_namedtuple(self):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)
diff --git a/psutil/tests/test_osx.py b/psutil/tests/test_osx.py
index 402571909..232b5a716 100755
--- a/psutil/tests/test_osx.py
+++ b/psutil/tests/test_osx.py
@@ -18,7 +18,6 @@
 from psutil.tests import MEMORY_TOLERANCE
 from psutil.tests import reap_children
 from psutil.tests import retry_before_failing
-from psutil.tests import run_test_module_by_name
 from psutil.tests import sh
 from psutil.tests import unittest
 
@@ -291,4 +290,5 @@ def test_sensors_battery(self):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)
diff --git a/psutil/tests/test_posix.py b/psutil/tests/test_posix.py
index 79931283d..88506b9af 100755
--- a/psutil/tests/test_posix.py
+++ b/psutil/tests/test_posix.py
@@ -30,7 +30,6 @@
 from psutil.tests import PYTHON_EXE
 from psutil.tests import reap_children
 from psutil.tests import retry_before_failing
-from psutil.tests import run_test_module_by_name
 from psutil.tests import sh
 from psutil.tests import skip_on_access_denied
 from psutil.tests import TRAVIS
@@ -437,4 +436,5 @@ def df(device):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)
diff --git a/psutil/tests/test_process.py b/psutil/tests/test_process.py
index 2af676ba9..a4d49ed4c 100755
--- a/psutil/tests/test_process.py
+++ b/psutil/tests/test_process.py
@@ -57,7 +57,6 @@
 from psutil.tests import PYTHON_EXE
 from psutil.tests import reap_children
 from psutil.tests import retry_before_failing
-from psutil.tests import run_test_module_by_name
 from psutil.tests import safe_rmpath
 from psutil.tests import sh
 from psutil.tests import skip_on_access_denied
@@ -1608,4 +1607,5 @@ def test_kill_terminate(self):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)
diff --git a/psutil/tests/test_sunos.py b/psutil/tests/test_sunos.py
index ea9afcde0..75b99feec 100755
--- a/psutil/tests/test_sunos.py
+++ b/psutil/tests/test_sunos.py
@@ -10,7 +10,6 @@
 
 import psutil
 from psutil import SUNOS
-from psutil.tests import run_test_module_by_name
 from psutil.tests import sh
 from psutil.tests import unittest
 
@@ -42,4 +41,5 @@ def test_cpu_count(self):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)
diff --git a/psutil/tests/test_system.py b/psutil/tests/test_system.py
index 51223fb00..43de8b99f 100755
--- a/psutil/tests/test_system.py
+++ b/psutil/tests/test_system.py
@@ -45,7 +45,6 @@
 from psutil.tests import mock
 from psutil.tests import reap_children
 from psutil.tests import retry_before_failing
-from psutil.tests import run_test_module_by_name
 from psutil.tests import safe_rmpath
 from psutil.tests import TESTFN
 from psutil.tests import TESTFN_UNICODE
@@ -868,4 +867,5 @@ def test_sensors_fans(self):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)
diff --git a/psutil/tests/test_unicode.py b/psutil/tests/test_unicode.py
index cfa8dd92c..a015c8961 100755
--- a/psutil/tests/test_unicode.py
+++ b/psutil/tests/test_unicode.py
@@ -77,7 +77,6 @@
 from psutil.tests import mock
 from psutil.tests import PYPY
 from psutil.tests import reap_children
-from psutil.tests import run_test_module_by_name
 from psutil.tests import safe_mkdir
 from psutil.tests import safe_rmpath as _safe_rmpath
 from psutil.tests import skip_on_access_denied
@@ -367,4 +366,5 @@ def test_proc_environ(self):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)
diff --git a/psutil/tests/test_windows.py b/psutil/tests/test_windows.py
index 222d8636f..00d0fece6 100755
--- a/psutil/tests/test_windows.py
+++ b/psutil/tests/test_windows.py
@@ -27,7 +27,6 @@
 from psutil.tests import mock
 from psutil.tests import reap_children
 from psutil.tests import retry_before_failing
-from psutil.tests import run_test_module_by_name
 from psutil.tests import sh
 from psutil.tests import unittest
 
@@ -865,4 +864,5 @@ def test_win_service_get(self):
 
 
 if __name__ == '__main__':
+    from psutil.tests.runner import run_test_module_by_name
     run_test_module_by_name(__file__)