terminal: add unittests (and minor tweaks) am: 788c6cb2ef

Original change: https://android-review.googlesource.com/c/platform/tools/repohooks/+/2620893

Change-Id: Ic84bc16963abf989a1797677be4a9b46f421d1c3
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

3 files changed, 175 insertions, 3 deletions

diff --git a/PREUPLOAD.cfg b/PREUPLOAD.cfg
index 631915d..8ee46ac 100644
--- a/PREUPLOAD.cfg
+++ b/PREUPLOAD.cfg
@@ -3,6 +3,7 @@
 config_unittest = ./rh/config_unittest.py
 hooks_unittest  = ./rh/hooks_unittest.py
 shell_unittest  = ./rh/shell_unittest.py
+terminal_unittest  = ./rh/terminal_unittest.py
 utils_unittest  = ./rh/utils_unittest.py
 android_test_mapping_format_unittest = ./tools/android_test_mapping_format_unittest.py
 clang-format unittest = ./tools/clang-format_unittest.py
diff --git a/rh/terminal.py b/rh/terminal.py
index feef884..c699fa4 100644
--- a/rh/terminal.py
+++ b/rh/terminal.py
@@ -29,6 +29,12 @@ del _path
 import rh.shell
 
 
+# This will erase all content in the current line after the cursor.  This is
+# useful for partial updates & progress messages as the terminal can display
+# it better.
+CSI_ERASE_LINE_AFTER = '\x1b[K'
+
+
 class Color(object):
     """Conditionally wraps text in ANSI color escape sequences."""
 
@@ -36,7 +42,7 @@ class Color(object):
     BOLD = -1
     COLOR_START = '\033[1;%dm'
     BOLD_START = '\033[1m'
-    RESET = '\033[0m'
+    RESET = '\033[m'
 
     def __init__(self, enabled=None):
         """Create a new Color object, optionally disabling color output.
@@ -51,7 +57,7 @@ class Color(object):
         """Returns a start color code.
 
         Args:
-          color: Color to use, .e.g BLACK, RED, etc.
+          color: Color to use, e.g. BLACK, RED, etc...
 
         Returns:
           If color is enabled, returns an ANSI sequence to start the given
@@ -111,7 +117,7 @@ def print_status_line(line, print_newline=False):
       print_newline: Print a newline at the end, if sys.stderr is a TTY.
     """
     if sys.stderr.isatty():
-        output = '\r' + line + '\x1B[K'
+        output = '\r' + line + CSI_ERASE_LINE_AFTER
         if print_newline:
             output += '\n'
     else:
diff --git a/rh/terminal_unittest.py b/rh/terminal_unittest.py
new file mode 100755
index 0000000..2239e7f
--- /dev/null
+++ b/rh/terminal_unittest.py
@@ -0,0 +1,165 @@
+#!/usr/bin/env python3
+# Copyright 2023 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unittests for the terminal module."""
+
+import contextlib
+import io
+import os
+import sys
+import unittest
+
+_path = os.path.realpath(__file__ + '/../..')
+if sys.path[0] != _path:
+    sys.path.insert(0, _path)
+del _path
+
+# We have to import our local modules after the sys.path tweak.  We can't use
+# relative imports because this is an executable program, not a module.
+# pylint: disable=wrong-import-position
+import rh.terminal
+
+
+class ColorTests(unittest.TestCase):
+    """Verify behavior of Color class."""
+
+    def setUp(self):
+        os.environ.pop('NOCOLOR', None)
+
+    def test_enabled_auto_tty(self):
+        """Test automatic enable behavior based on tty."""
+        stderr = io.StringIO()
+        with contextlib.redirect_stderr(stderr):
+            c = rh.terminal.Color()
+            self.assertFalse(c.enabled)
+
+            stderr.isatty = lambda: True
+            c = rh.terminal.Color()
+            self.assertTrue(c.enabled)
+
+    def test_enabled_auto_env(self):
+        """Test automatic enable behavior based on $NOCOLOR."""
+        stderr = io.StringIO()
+        with contextlib.redirect_stderr(stderr):
+            os.environ['NOCOLOR'] = 'yes'
+            c = rh.terminal.Color()
+            self.assertFalse(c.enabled)
+
+            os.environ['NOCOLOR'] = 'no'
+            c = rh.terminal.Color()
+            self.assertTrue(c.enabled)
+
+    def test_enabled_override(self):
+        """Test explicit enable behavior."""
+        stderr = io.StringIO()
+        with contextlib.redirect_stderr(stderr):
+            stderr.isatty = lambda: True
+            os.environ['NOCOLOR'] = 'no'
+            c = rh.terminal.Color()
+            self.assertTrue(c.enabled)
+            c = rh.terminal.Color(False)
+            self.assertFalse(c.enabled)
+
+            stderr.isatty = lambda: False
+            os.environ['NOCOLOR'] = 'yes'
+            c = rh.terminal.Color()
+            self.assertFalse(c.enabled)
+            c = rh.terminal.Color(True)
+            self.assertTrue(c.enabled)
+
+    def test_output_disabled(self):
+        """Test output when coloring is disabled."""
+        c = rh.terminal.Color(False)
+        self.assertEqual(c.start(rh.terminal.Color.BLACK), '')
+        self.assertEqual(c.color(rh.terminal.Color.BLACK, 'foo'), 'foo')
+        self.assertEqual(c.stop(), '')
+
+    def test_output_enabled(self):
+        """Test output when coloring is enabled."""
+        c = rh.terminal.Color(True)
+        self.assertEqual(c.start(rh.terminal.Color.BLACK), '\x1b[1;30m')
+        self.assertEqual(c.color(rh.terminal.Color.BLACK, 'foo'),
+                         '\x1b[1;30mfoo\x1b[m')
+        self.assertEqual(c.stop(), '\x1b[m')
+
+
+class PrintStatusLine(unittest.TestCase):
+    """Verify behavior of print_status_line."""
+
+    def test_terminal(self):
+        """Check tty behavior."""
+        stderr = io.StringIO()
+        stderr.isatty = lambda: True
+        with contextlib.redirect_stderr(stderr):
+            rh.terminal.print_status_line('foo')
+            rh.terminal.print_status_line('bar', print_newline=True)
+        csi = rh.terminal.CSI_ERASE_LINE_AFTER
+        self.assertEqual(stderr.getvalue(), f'\rfoo{csi}\rbar{csi}\n')
+
+    def test_no_terminal(self):
+        """Check tty-less behavior."""
+        stderr = io.StringIO()
+        with contextlib.redirect_stderr(stderr):
+            rh.terminal.print_status_line('foo')
+            rh.terminal.print_status_line('bar', print_newline=True)
+        self.assertEqual(stderr.getvalue(), 'foo\nbar\n')
+
+
+@contextlib.contextmanager
+def redirect_stdin(new_target):
+    """Temporarily switch sys.stdin to |new_target|."""
+    old = sys.stdin
+    try:
+        sys.stdin = new_target
+        yield
+    finally:
+        sys.stdin = old
+
+
+class BooleanPromptTests(unittest.TestCase):
+    """Verify behavior of boolean_prompt."""
+
+    def setUp(self):
+        self.stdin = io.StringIO()
+
+    def set_stdin(self, value: str) -> None:
+        """Set stdin wrapper to a string."""
+        self.stdin.seek(0)
+        self.stdin.write(value)
+        self.stdin.truncate()
+        self.stdin.seek(0)
+
+    def test_defaults(self):
+        """Test default behavior."""
+        stdout = io.StringIO()
+        with redirect_stdin(self.stdin), contextlib.redirect_stdout(stdout):
+            # Default values.  Will loop to EOF when it doesn't match anything.
+            for v in ('', '\n', 'oops'):
+                self.set_stdin(v)
+                self.assertTrue(rh.terminal.boolean_prompt())
+
+            # False values.
+            for v in ('n', 'N', 'no', 'NO'):
+                self.set_stdin(v)
+                self.assertFalse(rh.terminal.boolean_prompt())
+
+            # True values.
+            for v in ('y', 'Y', 'ye', 'yes', 'YES'):
+                self.set_stdin(v)
+                self.assertTrue(rh.terminal.boolean_prompt())
+
+
+if __name__ == '__main__':
+    unittest.main()