Starting on a basic colored output test runner using colorama. Haven't tested it on Win32 yet.

coderanger · coderanger · commit 1a8b904279d8 · 2010-08-09T05:13:45.000-07:00
diff --git a/pyzen/flaskext/script.py b/pyzen/flaskext/script.py
@@ -7,6 +7,7 @@
 from flaskext.script import Command, Option
 
 from pyzen.core import main
+from pyzen.runner import ColoredTextTestRunner
 
 try:
     from unittest2 import TestLoader
@@ -26,17 +27,21 @@ def _match_path(self, path, full_path, pattern):
                 return True
         return False
 
-def run_tests(app, pattern, start_dir, verbosity):
+def run_tests(app, pattern, start_dir, verbosity, nocolor):
     print start_dir
     loader = ZenTestLoader()
     suite = loader.discover(start_dir, pattern, start_dir)
-    result = unittest.TextTestRunner(verbosity=verbosity).run(suite)
+    if nocolor:
+        runner = unittest.TextTestRunner
+    else:
+        runner = ColoredTextTestRunner
+    result = runner(verbosity=verbosity).run(suite)
     return result
 
 class Test(Command):
     """Run app tests."""
     
-    def __init__(self, pattern='*/tests/*.py;*/tests.py', start_dir=None, verbosity=0):
+    def __init__(self, pattern='*/tests/*.py;*/tests.py', start_dir=None, verbosity=1):
         if start_dir is None:
             # Find the file that called this constructor and use its directory
             for f in inspect.stack():
@@ -54,10 +59,11 @@ def get_options(self):
             Option('-p', '--pattern', dest='pattern', default=self.default_pattern),
             Option('-s', '--start_dir', dest='start_dir', default=self.default_start_dir),
             Option('-v', '--verbosity', dest='verbosity', default=self.default_verbosity),
+            Option('--nocolor', action='store_true', default=False, help='Disable colored output'),
         ]
     
-    def run(self, app, pattern, start_dir, verbosity):
-        result = run_tests(app, pattern, start_dir, verbosity)
+    def run(self, app, pattern, start_dir, verbosity, nocolor):
+        result = run_tests(app, pattern, start_dir, verbosity, nocolor)
         if result.failures or result.errors:
             sys.exit(1)
 
@@ -70,8 +76,8 @@ def get_options(self):
         options.append(Option('-u', '--ui', help='Force the use of the given PyZen UI'))
         return options
     
-    def run(self, app, pattern, start_dir, verbosity, ui):
+    def run(self, app, pattern, start_dir, verbosity, ui, nocolor):
         try:
-            main(ui, run_tests, app, pattern, start_dir, verbosity)
+            main(ui, run_tests, app, pattern, start_dir, verbosity, nocolor)
         except KeyboardInterrupt:
             pass
diff --git a/pyzen/runner.py b/pyzen/runner.py
@@ -0,0 +1,73 @@
+import time
+import unittest
+
+try:
+    import colorama
+except ImportError:
+    from pyzen import _colorama as colorama
+
+COLOR_SUCCESS = colorama.Fore.GREEN
+COLOR_FAIL = colorama.Fore.RED
+COLOR_RESET = colorama.Fore.RESET
+
+class ColoredTextTestResult(unittest._TextTestResult):
+    
+    def addSuccess(self, test):
+        self.stream.write(COLOR_SUCCESS)
+        unittest._TextTestResult.addSuccess(self, test)
+        self.stream.write(COLOR_RESET)
+    
+    def addError(self, test, err):
+        self.stream.write(COLOR_FAIL)
+        unittest._TextTestResult.addError(self, test, err)
+        self.stream.write(COLOR_RESET)
+    
+    def addFailure(self, test, err):
+        self.stream.write(COLOR_FAIL)
+        unittest._TextTestResult.addFailure(self, test, err)
+        self.stream.write(COLOR_RESET)
+
+    def printErrorList(self, flavour, errors):
+            for test, err in errors:
+                self.stream.writeln(self.separator1)
+                self.stream.writeln("%s%s: %s%s" % (COLOR_FAIL,flavour,self.getDescription(test),COLOR_RESET))
+                self.stream.writeln(self.separator2)
+                self.stream.writeln("%s" % err)
+
+class ColoredTextTestRunner(unittest.TextTestRunner):
+    
+    def __init__(self, *args, **kwargs):
+        unittest.TextTestRunner.__init__(self, *args, **kwargs)
+        wrapper = colorama.AnsiToWin32(self.stream)
+        if wrapper.should_wrap():
+            self.stream = wrapper.stream
+    
+    def _makeResult(self):
+        return ColoredTextTestResult(self.stream, self.descriptions, self.verbosity)
+    
+    def run(self, test):
+        "Run the given test case or test suite."
+        result = self._makeResult()
+        startTime = time.time()
+        test(result)
+        stopTime = time.time()
+        timeTaken = stopTime - startTime
+        result.printErrors()
+        self.stream.writeln(result.separator2)
+        run = result.testsRun
+        self.stream.writeln("Ran %d test%s in %.3fs" %
+                            (run, run != 1 and "s" or "", timeTaken))
+        self.stream.writeln()
+        if not result.wasSuccessful():
+            self.stream.write(COLOR_FAIL)
+            self.stream.write("FAILED (")
+            failed, errored = map(len, (result.failures, result.errors))
+            if failed:
+                self.stream.write("failures=%d" % failed)
+            if errored:
+                if failed: self.stream.write(", ")
+                self.stream.write("errors=%d" % errored)
+            self.stream.writeln(")"+COLOR_RESET)
+        else:
+            self.stream.writeln(COLOR_SUCCESS+"OK"+COLOR_RESET)
+        return result
