1# Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Tests of the readline-based CLI.""" 16import os 17import argparse 18import tempfile 19 20from tensorflow.python.debug.cli import cli_config 21from tensorflow.python.debug.cli import debugger_cli_common 22from tensorflow.python.debug.cli import readline_ui 23from tensorflow.python.debug.cli import ui_factory 24from tensorflow.python.framework import test_util 25from tensorflow.python.lib.io import file_io 26from tensorflow.python.platform import gfile 27from tensorflow.python.platform import googletest 28 29 30class MockReadlineUI(readline_ui.ReadlineUI): 31 """Test subclass of ReadlineUI that bypasses terminal manipulations.""" 32 33 def __init__(self, on_ui_exit=None, command_sequence=None): 34 _, config_file_path = tempfile.mkstemp() # safe to ignore fd 35 readline_ui.ReadlineUI.__init__( 36 self, 37 on_ui_exit=on_ui_exit, 38 config=cli_config.CLIConfig(config_file_path=config_file_path)) 39 40 self._command_sequence = command_sequence 41 self._command_counter = 0 42 43 self.observers = {"screen_outputs": []} 44 45 def _get_user_command(self): 46 command = self._command_sequence[self._command_counter] 47 self._command_counter += 1 48 return command 49 50 def _display_output(self, screen_output): 51 self.observers["screen_outputs"].append(screen_output) 52 53 54class CursesTest(test_util.TensorFlowTestCase): 55 56 def setUp(self): 57 self._tmp_dir = tempfile.mkdtemp() 58 self._tmp_config_path = os.path.join(self._tmp_dir, ".tfdbg_config") 59 self.assertFalse(gfile.Exists(self._tmp_config_path)) 60 super(CursesTest, self).setUp() 61 62 def tearDown(self): 63 file_io.delete_recursively(self._tmp_dir) 64 super(CursesTest, self).tearDown() 65 66 def _babble(self, args, screen_info=None): 67 ap = argparse.ArgumentParser( 68 description="Do babble.", usage=argparse.SUPPRESS) 69 ap.add_argument( 70 "-n", 71 "--num_times", 72 dest="num_times", 73 type=int, 74 default=60, 75 help="How many times to babble") 76 77 parsed = ap.parse_args(args) 78 79 lines = ["bar"] * parsed.num_times 80 return debugger_cli_common.RichTextLines(lines) 81 82 def testUIFactoryCreatesReadlineUI(self): 83 ui = ui_factory.get_ui( 84 "readline", 85 config=cli_config.CLIConfig(config_file_path=self._tmp_config_path)) 86 self.assertIsInstance(ui, readline_ui.ReadlineUI) 87 88 def testUIFactoryRaisesExceptionOnInvalidUIType(self): 89 with self.assertRaisesRegex(ValueError, "Invalid ui_type: 'foobar'"): 90 ui_factory.get_ui( 91 "foobar", 92 config=cli_config.CLIConfig(config_file_path=self._tmp_config_path)) 93 94 def testUIFactoryRaisesExceptionOnInvalidUITypeGivenAvailable(self): 95 with self.assertRaisesRegex(ValueError, "Invalid ui_type: 'readline'"): 96 ui_factory.get_ui( 97 "readline", 98 available_ui_types=["curses"], 99 config=cli_config.CLIConfig(config_file_path=self._tmp_config_path)) 100 101 def testRunUIExitImmediately(self): 102 """Make sure that the UI can exit properly after launch.""" 103 104 ui = MockReadlineUI(command_sequence=["exit"]) 105 ui.run_ui() 106 107 # No screen output should have happened. 108 self.assertEqual(0, len(ui.observers["screen_outputs"])) 109 110 def testRunUIEmptyCommand(self): 111 """Issue an empty command then exit.""" 112 113 ui = MockReadlineUI(command_sequence=["", "exit"]) 114 ui.run_ui() 115 self.assertEqual(1, len(ui.observers["screen_outputs"])) 116 117 def testRunUIWithInitCmd(self): 118 """Run UI with an initial command specified.""" 119 120 ui = MockReadlineUI(command_sequence=["exit"]) 121 122 ui.register_command_handler("babble", self._babble, "") 123 ui.run_ui(init_command="babble") 124 125 screen_outputs = ui.observers["screen_outputs"] 126 self.assertEqual(1, len(screen_outputs)) 127 self.assertEqual(["bar"] * 60, screen_outputs[0].lines) 128 129 def testRunUIWithValidUsersCommands(self): 130 """Run UI with an initial command specified.""" 131 132 ui = MockReadlineUI(command_sequence=["babble -n 3", "babble -n 6", "exit"]) 133 ui.register_command_handler("babble", self._babble, "") 134 ui.run_ui() 135 136 screen_outputs = ui.observers["screen_outputs"] 137 self.assertEqual(2, len(screen_outputs)) 138 self.assertEqual(["bar"] * 3, screen_outputs[0].lines) 139 self.assertEqual(["bar"] * 6, screen_outputs[1].lines) 140 141 def testRunUIWithInvalidUsersCommands(self): 142 """Run UI with an initial command specified.""" 143 144 ui = MockReadlineUI(command_sequence=["babble -n 3", "wobble", "exit"]) 145 ui.register_command_handler("babble", self._babble, "") 146 ui.run_ui() 147 148 screen_outputs = ui.observers["screen_outputs"] 149 self.assertEqual(2, len(screen_outputs)) 150 self.assertEqual(["bar"] * 3, screen_outputs[0].lines) 151 self.assertEqual(["ERROR: Invalid command prefix \"wobble\""], 152 screen_outputs[1].lines) 153 154 def testRunUIWithOnUIExitCallback(self): 155 observer = {"callback_invoked": False} 156 157 def callback_for_test(): 158 observer["callback_invoked"] = True 159 160 ui = MockReadlineUI(on_ui_exit=callback_for_test, command_sequence=["exit"]) 161 162 self.assertFalse(observer["callback_invoked"]) 163 ui.run_ui() 164 165 self.assertEqual(0, len(ui.observers["screen_outputs"])) 166 self.assertTrue(observer["callback_invoked"]) 167 168 def testIncompleteRedirectWorks(self): 169 _, output_path = tempfile.mkstemp() # safe to ignore fd 170 171 ui = MockReadlineUI( 172 command_sequence=["babble -n 2 > %s" % output_path, "exit"]) 173 174 ui.register_command_handler("babble", self._babble, "") 175 ui.run_ui() 176 177 screen_outputs = ui.observers["screen_outputs"] 178 self.assertEqual(1, len(screen_outputs)) 179 self.assertEqual(["bar"] * 2, screen_outputs[0].lines) 180 181 with gfile.Open(output_path, "r") as f: 182 self.assertEqual("bar\nbar\n", f.read()) 183 184 def testConfigSetAndShow(self): 185 """Run UI with an initial command specified.""" 186 187 ui = MockReadlineUI(command_sequence=[ 188 "config set graph_recursion_depth 5", "config show", "exit"]) 189 ui.run_ui() 190 outputs = ui.observers["screen_outputs"] 191 self.assertEqual( 192 ["Command-line configuration:", 193 "", 194 " graph_recursion_depth: 5"], outputs[1].lines[:3]) 195 196 197if __name__ == "__main__": 198 googletest.main() 199