rh/terminal_unittest.py

#!/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 StringPromptTests(unittest.TestCase):
    """Verify behavior of str_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):
            # Test EOF behavior.
            self.assertIsNone(rh.terminal.str_prompt('foo', ('a', 'b')))

            # Test enter behavior.
            self.set_stdin('\n')
            self.assertEqual(rh.terminal.str_prompt('foo', ('a', 'b')), '')

            # Lowercase inputs.
            self.set_stdin('Ok')
            self.assertEqual(rh.terminal.str_prompt('foo', ('a', 'b')), 'ok')

            # Don't lowercase inputs.
            self.set_stdin('Ok')
            self.assertEqual(
                rh.terminal.str_prompt('foo', ('a', 'b'), lower=False), 'Ok')


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()