1#!/usr/bin/env python3 2# Copyright 2023 The Android Open Source Project 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15 16"""Unittests for the terminal module.""" 17 18import contextlib 19import io 20import os 21import sys 22import unittest 23 24_path = os.path.realpath(__file__ + '/../..') 25if sys.path[0] != _path: 26 sys.path.insert(0, _path) 27del _path 28 29# We have to import our local modules after the sys.path tweak. We can't use 30# relative imports because this is an executable program, not a module. 31# pylint: disable=wrong-import-position 32import rh.terminal 33 34 35class ColorTests(unittest.TestCase): 36 """Verify behavior of Color class.""" 37 38 def setUp(self): 39 os.environ.pop('NOCOLOR', None) 40 41 def test_enabled_auto_tty(self): 42 """Test automatic enable behavior based on tty.""" 43 stderr = io.StringIO() 44 with contextlib.redirect_stderr(stderr): 45 c = rh.terminal.Color() 46 self.assertFalse(c.enabled) 47 48 stderr.isatty = lambda: True 49 c = rh.terminal.Color() 50 self.assertTrue(c.enabled) 51 52 def test_enabled_auto_env(self): 53 """Test automatic enable behavior based on $NOCOLOR.""" 54 stderr = io.StringIO() 55 with contextlib.redirect_stderr(stderr): 56 os.environ['NOCOLOR'] = 'yes' 57 c = rh.terminal.Color() 58 self.assertFalse(c.enabled) 59 60 os.environ['NOCOLOR'] = 'no' 61 c = rh.terminal.Color() 62 self.assertTrue(c.enabled) 63 64 def test_enabled_override(self): 65 """Test explicit enable behavior.""" 66 stderr = io.StringIO() 67 with contextlib.redirect_stderr(stderr): 68 stderr.isatty = lambda: True 69 os.environ['NOCOLOR'] = 'no' 70 c = rh.terminal.Color() 71 self.assertTrue(c.enabled) 72 c = rh.terminal.Color(False) 73 self.assertFalse(c.enabled) 74 75 stderr.isatty = lambda: False 76 os.environ['NOCOLOR'] = 'yes' 77 c = rh.terminal.Color() 78 self.assertFalse(c.enabled) 79 c = rh.terminal.Color(True) 80 self.assertTrue(c.enabled) 81 82 def test_output_disabled(self): 83 """Test output when coloring is disabled.""" 84 c = rh.terminal.Color(False) 85 self.assertEqual(c.start(rh.terminal.Color.BLACK), '') 86 self.assertEqual(c.color(rh.terminal.Color.BLACK, 'foo'), 'foo') 87 self.assertEqual(c.stop(), '') 88 89 def test_output_enabled(self): 90 """Test output when coloring is enabled.""" 91 c = rh.terminal.Color(True) 92 self.assertEqual(c.start(rh.terminal.Color.BLACK), '\x1b[1;30m') 93 self.assertEqual(c.color(rh.terminal.Color.BLACK, 'foo'), 94 '\x1b[1;30mfoo\x1b[m') 95 self.assertEqual(c.stop(), '\x1b[m') 96 97 98class PrintStatusLine(unittest.TestCase): 99 """Verify behavior of print_status_line.""" 100 101 def test_terminal(self): 102 """Check tty behavior.""" 103 stderr = io.StringIO() 104 stderr.isatty = lambda: True 105 with contextlib.redirect_stderr(stderr): 106 rh.terminal.print_status_line('foo') 107 rh.terminal.print_status_line('bar', print_newline=True) 108 csi = rh.terminal.CSI_ERASE_LINE_AFTER 109 self.assertEqual(stderr.getvalue(), f'\rfoo{csi}\rbar{csi}\n') 110 111 def test_no_terminal(self): 112 """Check tty-less behavior.""" 113 stderr = io.StringIO() 114 with contextlib.redirect_stderr(stderr): 115 rh.terminal.print_status_line('foo') 116 rh.terminal.print_status_line('bar', print_newline=True) 117 self.assertEqual(stderr.getvalue(), 'foo\nbar\n') 118 119 120@contextlib.contextmanager 121def redirect_stdin(new_target): 122 """Temporarily switch sys.stdin to |new_target|.""" 123 old = sys.stdin 124 try: 125 sys.stdin = new_target 126 yield 127 finally: 128 sys.stdin = old 129 130 131class StringPromptTests(unittest.TestCase): 132 """Verify behavior of str_prompt.""" 133 134 def setUp(self): 135 self.stdin = io.StringIO() 136 137 def set_stdin(self, value: str) -> None: 138 """Set stdin wrapper to a string.""" 139 self.stdin.seek(0) 140 self.stdin.write(value) 141 self.stdin.truncate() 142 self.stdin.seek(0) 143 144 def test_defaults(self): 145 """Test default behavior.""" 146 stdout = io.StringIO() 147 with redirect_stdin(self.stdin), contextlib.redirect_stdout(stdout): 148 # Test EOF behavior. 149 self.assertIsNone(rh.terminal.str_prompt('foo', ('a', 'b'))) 150 151 # Test enter behavior. 152 self.set_stdin('\n') 153 self.assertEqual(rh.terminal.str_prompt('foo', ('a', 'b')), '') 154 155 # Lowercase inputs. 156 self.set_stdin('Ok') 157 self.assertEqual(rh.terminal.str_prompt('foo', ('a', 'b')), 'ok') 158 159 # Don't lowercase inputs. 160 self.set_stdin('Ok') 161 self.assertEqual( 162 rh.terminal.str_prompt('foo', ('a', 'b'), lower=False), 'Ok') 163 164 165class BooleanPromptTests(unittest.TestCase): 166 """Verify behavior of boolean_prompt.""" 167 168 def setUp(self): 169 self.stdin = io.StringIO() 170 171 def set_stdin(self, value: str) -> None: 172 """Set stdin wrapper to a string.""" 173 self.stdin.seek(0) 174 self.stdin.write(value) 175 self.stdin.truncate() 176 self.stdin.seek(0) 177 178 def test_defaults(self): 179 """Test default behavior.""" 180 stdout = io.StringIO() 181 with redirect_stdin(self.stdin), contextlib.redirect_stdout(stdout): 182 # Default values. Will loop to EOF when it doesn't match anything. 183 for v in ('', '\n', 'oops'): 184 self.set_stdin(v) 185 self.assertTrue(rh.terminal.boolean_prompt()) 186 187 # False values. 188 for v in ('n', 'N', 'no', 'NO'): 189 self.set_stdin(v) 190 self.assertFalse(rh.terminal.boolean_prompt()) 191 192 # True values. 193 for v in ('y', 'Y', 'ye', 'yes', 'YES'): 194 self.set_stdin(v) 195 self.assertTrue(rh.terminal.boolean_prompt()) 196 197 198if __name__ == '__main__': 199 unittest.main() 200