mavsdk_tests: import naming, easier coloring

test/mavsdk_tests/logger_helper.py

@@ -3,50 +3,52 @@
 import re
 import sys
 import os
+from enum import Enum
+from functools import lru_cache
 
 
-def supports_color():
+class color(Enum):
-    """Returns True if the running system's terminal supports color.
+    PURPLE = '\033[95m'
-
+    CYAN = '\033[96m'
-    From https://stackoverflow.com/a/22254892/8548472
+    DARKCYAN = '\033[36m'
-    """
+    BLUE = '\033[94m'
-    plat = sys.platform
+    GREEN = '\033[92m'
-    supported_platform = plat != 'Pocket PC' and (plat != 'win32' or
+    YELLOW = '\033[93m'
-                                                  'ANSICON' in os.environ)
+    RED = '\033[91m'
-    # isatty is not always implemented, #6223.
+    BOLD = '\033[1m'
-    is_a_tty = hasattr(sys.stdout, 'isatty') and sys.stdout.isatty()
+    UNDERLINE = '\033[4m'
-    return supported_platform and is_a_tty
+    RESET = '\033[0m'
 
 
-if supports_color():
+def colorize(text: str, c: color) -> str:
-    class color:
+    if _supports_color():
-        PURPLE = '\033[95m'
+        return c.value + text + color.RESET.value
-        CYAN = '\033[96m'
+    else:
-        DARKCYAN = '\033[36m'
+        return text
-        BLUE = '\033[94m'
-        GREEN = '\033[92m'
-        YELLOW = '\033[93m'
-        RED = '\033[91m'
-        BOLD = '\033[1m'
-        UNDERLINE = '\033[4m'
-        END = '\033[0m'
-else:
-    class color:
-        PURPLE = ''
-        CYAN = ''
-        DARKCYAN = ''
-        BLUE = ''
-        GREEN = ''
-        YELLOW = ''
-        RED = ''
-        BOLD = ''
-        UNDERLINE = ''
-        END = ''
 
 
-def remove_color(text):
+def maybe_strip_color(text: str) -> str:
     """Remove ANSI and xterm256 color codes.
 
     From https://stackoverflow.com/a/30500866/8548472
     """
-    return re.sub(r'\x1b(\[.*?[@-~]|\].*?(\x07|\x1b\\))', '', text)
+    if not _supports_color():
+        return re.sub(r'\x1b(\[.*?[@-~]|\].*?(\x07|\x1b\\))', '', text)
+    else:
+        return text
+
+
+@lru_cache
+def _supports_color() -> bool:
+    """Returns True if the running system's terminal supports color.
+
+    From https://stackoverflow.com/a/22254892/8548472
+    """
+    supported_platform = \
+        (sys.platform != 'Pocket PC') and \
+        (sys.platform != 'win32' or 'ANSICON' in os.environ)
+
+    # isatty is not always implemented.
+    is_a_tty = hasattr(sys.stdout, 'isatty') and sys.stdout.isatty()
+
+    return supported_platform and is_a_tty

test/mavsdk_tests/mavsdk_test_runner.py

@@ -7,12 +7,8 @@ import psutil
 import signal
 import subprocess
 import sys
-from logger_helper import color
+from logger_helper import color, colorize
-from process_helper import \
+import process_helper as ph
-    Px4Runner, \
-    GzserverRunner, \
-    GzclientRunner, \
-    TestRunner
 
 
 def main():
@@ -160,10 +156,12 @@ class Tester:
                 break
 
         if overall_success:
-            print(color.GREEN + "Overall result: success!" + color.END)
+            print(colorize(colorize(
+                "Overall result: success!", color.GREEN), color.BOLD))
             return 0
         else:
-            print(color.RED + "Overall result: failure!" + color.END)
+            print(colorize(colorize(
+                "Overall result: failure!", color.RED), color.BOLD))
             return 1
 
     def run_test_group(self):
@@ -184,8 +182,10 @@ class Tester:
                 models = []
 
         for group in models:
-            print(color.BOLD + "==> Running tests for '{}' with filter '{}'"
+            print(colorize(
-                  .format(group['model'], group['test_filter']) + color.END)
+                "==> Running tests for '{}' with filter '{}'"
+                .format(group['model'], group['test_filter']),
+                color.BOLD))
 
             tests = determine_tests(group['test_filter'])
 
@@ -194,11 +194,15 @@ class Tester:
                       format(i+1, len(tests), test))
                 was_success = self.run_test(test, group)
                 if was_success:
-                    print(color.GREEN + "--- Test {} of {}: '{}' succeeded.".
+                    print(colorize(
-                          format(i+1, len(tests), test) + color.END)
+                        "--- Test {} of {}: '{}' succeeded."
+                        .format(i+1, len(tests), test),
+                        color.GREEN))
                 else:
-                    print(color.RED + "--- Test {} of {}: '{}' failed.".
+                    print(colorize(
-                          format(i+1, len(tests), test) + color.END)
+                        "--- Test {} of {}: '{}' failed."
+                        .format(i+1, len(tests), test),
+                        color.RED))
 
                 if not was_success:
                     overall_success = False
@@ -217,14 +221,14 @@ class Tester:
             speed_factor = min(int(speed_factor), group["max_speed_factor"])
 
         if self.config['mode'] == 'sitl':
-            px4_runner = Px4Runner(
+            px4_runner = ph.Px4Runner(
                 group['model'], os.getcwd(), self.log_dir, speed_factor,
                 self.debugger, self.verbose)
             px4_runner.start(group)
             self.active_runners.append(px4_runner)
 
             if self.config['simulator'] == 'gazebo':
-                gzserver_runner = GzserverRunner(
+                gzserver_runner = ph.GzserverRunner(
                     group['model'],
                     os.getcwd(),
                     self.log_dir,
@@ -234,12 +238,12 @@ class Tester:
                 self.active_runners.append(gzserver_runner)
 
                 if self.gui:
-                    gzclient_runner = GzclientRunner(
+                    gzclient_runner = ph.GzclientRunner(
                         os.getcwd(), self.log_dir, self.verbose)
                     gzclient_runner.start(group)
                     self.active_runners.append(gzclient_runner)
 
-        test_runner = TestRunner(
+        test_runner = ph.TestRunner(
             os.getcwd(),
             self.log_dir,
             group,
@@ -261,8 +265,10 @@ class Tester:
                     runner.print_output()
 
         else:
-            print(color.BOLD + "Test timeout of {} mins triggered!".
+            print(colorize(
-                  format(group['timeout_min']) + color.END)
+                "Test timeout of {} mins triggered!".
+                format(group['timeout_min']),
+                color.BOLD))
             is_success = False
 
         for runner in self.active_runners:

test/mavsdk_tests/process_helper.py

@@ -8,7 +8,7 @@ import atexit
 import subprocess
 import threading
 import errno
-from logger_helper import color
+from logger_helper import color, maybe_strip_color, colorize
 
 
 class Runner:
@@ -81,10 +81,7 @@ class Runner:
     def get_output(self):
         try:
             output = self.output_queue.get(block=True, timeout=0.1)
-            if supports_color():
+            return maybe_strip_color(output)
-                return output
-            else:
-                return remove_color(output)
         except queue.Empty:
             return None
 
@@ -92,10 +89,8 @@ class Runner:
         output = self.get_output()
         if not output:
             return
-        print(color.END +
+        print(colorize("[" + self.name.ljust(11) + "] " + output, color.RESET),
-              "[" + self.name.ljust(11) + "] " +
+              end="")
-              output +
-              color.END, end="")
 
     def stop(self):
         atexit.unregister(self.stop)