1# 2# Copyright (C) 2024 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 17 18import builtins 19import io 20from io import BytesIO 21import socketserver 22import unittest 23import subprocess 24import sys 25import os 26import webbrowser 27from unittest import mock 28from src.command import OpenCommand 29from src.open_ui import download_trace_processor, open_trace 30 31ANDROID_BUILD_TOP = "/main" 32TEST_FILE = "file.pbtxt" 33TORQ_TEMP_DIR = "/tmp/.torq" 34PERFETTO_BINARY = "/trace_processor" 35TORQ_TEMP_TRACE_PROCESSOR = TORQ_TEMP_DIR + PERFETTO_BINARY 36ANDROID_PERFETTO_TOOLS_DIR = "/external/perfetto/tools" 37ANDROID_TRACE_PROCESSOR = ANDROID_PERFETTO_TOOLS_DIR + PERFETTO_BINARY 38LARGE_FILE_SIZE = 1024 * 1024 * 512 # 512 MB 39WEB_UI_ADDRESS = "https://ui.perfetto.dev" 40 41 42class OpenUiUnitTest(unittest.TestCase): 43 @mock.patch.dict(os.environ, {"ANDROID_BUILD_TOP": ANDROID_BUILD_TOP}, 44 clear=True) 45 @mock.patch.object(os.path, "exists", autospec=True) 46 @mock.patch.object(os.path, "getsize", autospec=True) 47 @mock.patch.object(subprocess, "run", autospec=True) 48 @mock.patch.object(builtins, "input") 49 def test_download_trace_processor_successfully(self, mock_input, 50 mock_subprocess_run, mock_getsize, mock_exists): 51 mock_getsize.return_value = LARGE_FILE_SIZE 52 mock_input.return_value = "y" 53 mock_exists.side_effect = [False, False, True] 54 mock_subprocess_run.return_value = None 55 terminal_output = io.StringIO() 56 sys.stdout = terminal_output 57 58 trace_processor_path = download_trace_processor(TEST_FILE) 59 60 self.assertEqual(trace_processor_path, TORQ_TEMP_TRACE_PROCESSOR) 61 self.assertEqual(terminal_output.getvalue(), "") 62 63 @mock.patch.dict(os.environ, {"ANDROID_BUILD_TOP": ANDROID_BUILD_TOP}, 64 clear=True) 65 @mock.patch.object(os.path, "exists", autospec=True) 66 @mock.patch.object(os.path, "getsize", autospec=True) 67 @mock.patch.object(subprocess, "run", autospec=True) 68 @mock.patch.object(builtins, "input") 69 def test_download_trace_processor_failed(self, mock_input, 70 mock_subprocess_run, mock_getsize, mock_exists): 71 mock_getsize.return_value = LARGE_FILE_SIZE 72 mock_input.return_value = "y" 73 mock_exists.side_effect = [False, False, False] 74 mock_subprocess_run.return_value = None 75 terminal_output = io.StringIO() 76 sys.stdout = terminal_output 77 78 trace_processor_path = download_trace_processor(TEST_FILE) 79 80 self.assertEqual(trace_processor_path, None) 81 self.assertEqual(terminal_output.getvalue(), 82 "Could not download perfetto scripts. Continuing.\n") 83 84 @mock.patch.dict(os.environ, {"ANDROID_BUILD_TOP": ANDROID_BUILD_TOP}, 85 clear=True) 86 @mock.patch.object(os.path, "exists", autospec=True) 87 @mock.patch.object(os.path, "getsize", autospec=True) 88 @mock.patch.object(builtins, "input") 89 def test_download_trace_processor_wrong_input(self, mock_input, 90 mock_getsize, mock_exists): 91 mock_getsize.return_value = LARGE_FILE_SIZE 92 mock_input.return_value = "bad-input" 93 mock_exists.side_effect = [False, False] 94 terminal_output = io.StringIO() 95 sys.stdout = terminal_output 96 97 error = download_trace_processor(TEST_FILE) 98 99 self.assertNotEqual(error, None) 100 self.assertEqual(error.message, "Invalid inputs.") 101 self.assertEqual(error.suggestion, "Please accept or reject the download.") 102 self.assertEqual(terminal_output.getvalue(), 103 "Invalid input. Please try again.\n" 104 "Invalid input. Please try again.\n") 105 106 @mock.patch.dict(os.environ, {"ANDROID_BUILD_TOP": ANDROID_BUILD_TOP}, 107 clear=True) 108 @mock.patch.object(os.path, "exists", autospec=True) 109 @mock.patch.object(os.path, "getsize", autospec=True) 110 @mock.patch.object(builtins, "input") 111 def test_download_trace_processor_refused(self, mock_input, mock_getsize, 112 mock_exists): 113 mock_getsize.return_value = LARGE_FILE_SIZE 114 mock_input.return_value = "n" 115 mock_exists.side_effect = [False, False] 116 terminal_output = io.StringIO() 117 sys.stdout = terminal_output 118 119 trace_processor_path = download_trace_processor(TEST_FILE) 120 121 self.assertEqual(trace_processor_path, None) 122 self.assertEqual(terminal_output.getvalue(), 123 "Will continue without downloading perfetto scripts.\n") 124 125 @mock.patch.dict(os.environ, {"ANDROID_BUILD_TOP": ANDROID_BUILD_TOP}, 126 clear=True) 127 @mock.patch.object(os.path, "getsize", autospec=True) 128 @mock.patch.object(os.path, "abspath", autospec=True) 129 @mock.patch.object(webbrowser, "open_new_tab", autospec=True) 130 @mock.patch.object(subprocess, "run", autospec=True) 131 @mock.patch.object(os.path, "exists", autospec=True) 132 @mock.patch.object(os.path, "expanduser", autospec=True) 133 @mock.patch.object(subprocess, "Popen", autospec=True) 134 def test_open_trace_scripts_large_file(self, mock_popen, mock_expanduser, 135 mock_exists, mock_subprocess_run, mock_open_new_tab, mock_abspath, 136 mock_getsize): 137 mock_expanduser.return_value = "" 138 mock_subprocess_run.return_value = None 139 mock_open_new_tab.return_value = None 140 mock_getsize.return_value = LARGE_FILE_SIZE 141 mock_abspath.return_value = TEST_FILE 142 mock_exists.return_value = True 143 terminal_output = io.StringIO() 144 sys.stdout = terminal_output 145 mock_process = mock_popen.return_value 146 mock_process.stdout = BytesIO(b'Trace loaded') 147 148 error = open_trace(TEST_FILE, WEB_UI_ADDRESS, False) 149 150 mock_open_new_tab.assert_called() 151 self.assertEqual(error, None) 152 self.assertEqual(terminal_output.getvalue(), 153 "\033[93m##### Loading trace. #####\n##### " 154 "Follow the directions in the Perfetto UI. Do not exit " 155 "out of torq until you are done viewing the trace. Press " 156 "CTRL+C to exit torq and close the trace_processor. " 157 "#####\033[0m\nProcess was killed.\n") 158 159 @mock.patch.dict(os.environ, {"ANDROID_BUILD_TOP": ANDROID_BUILD_TOP}, 160 clear=True) 161 @mock.patch.object(os.path, "exists", autospec=True) 162 @mock.patch.object(os.path, "getsize", autospec=True) 163 @mock.patch.object(webbrowser, "open_new_tab", autospec=True) 164 @mock.patch.object(subprocess, "run", autospec=True) 165 @mock.patch.object(subprocess, "Popen", autospec=True) 166 def test_open_trace_scripts_large_file_use_trace_processor_enabled(self, 167 mock_popen, mock_subprocess_run, mock_open_new_tab, mock_getsize, 168 mock_exists): 169 mock_subprocess_run.return_value = None 170 mock_open_new_tab.return_value = None 171 mock_getsize.return_value = LARGE_FILE_SIZE 172 mock_exists.return_value = True 173 terminal_output = io.StringIO() 174 sys.stdout = terminal_output 175 mock_process = mock_popen.return_value 176 mock_process.stdout = BytesIO(b'Trace loaded') 177 178 error = open_trace(TEST_FILE, WEB_UI_ADDRESS, True) 179 180 mock_open_new_tab.assert_called() 181 self.assertEqual(error, None) 182 self.assertEqual(terminal_output.getvalue(), 183 "\033[93m##### Loading trace. #####\n##### " 184 "Follow the directions in the Perfetto UI. Do not exit " 185 "out of torq until you are done viewing the trace. Press " 186 "CTRL+C to exit torq and close the trace_processor. " 187 "#####\033[0m\nProcess was killed.\n") 188 189 @mock.patch.object(os.path, "getsize", autospec=True) 190 @mock.patch.object(webbrowser, "open_new_tab", autospec=True) 191 @mock.patch.object(socketserver, "TCPServer", autospec=True) 192 def test_open_trace_scripts_small_file(self, mock_tcpserver, 193 mock_open_new_tab, mock_getsize): 194 def handle_request(): 195 mock_process.fname_get_completed = 0 196 return 197 mock_process = type('', (), {})() 198 mock_process.handle_request = handle_request 199 mock_tcpserver.return_value.__enter__.return_value = mock_process 200 mock_open_new_tab.return_value = None 201 mock_getsize.return_value = 0 202 terminal_output = io.StringIO() 203 sys.stdout = terminal_output 204 205 error = open_trace(TEST_FILE, WEB_UI_ADDRESS, False) 206 207 mock_open_new_tab.assert_called() 208 self.assertEqual(error, None) 209 self.assertEqual(terminal_output.getvalue(), "") 210 211 @mock.patch.dict(os.environ, {"ANDROID_BUILD_TOP": ANDROID_BUILD_TOP}, 212 clear=True) 213 @mock.patch.object(os.path, "exists", autospec=True) 214 @mock.patch.object(os.path, "getsize", autospec=True) 215 @mock.patch.object(webbrowser, "open_new_tab", autospec=True) 216 @mock.patch.object(subprocess, "run", autospec=True) 217 @mock.patch.object(subprocess, "Popen", autospec=True) 218 def test_download_trace_processor_small_file_use_trace_processor_enabled(self, 219 mock_popen, mock_subprocess_run, mock_open_new_tab, mock_getsize, 220 mock_exists): 221 mock_subprocess_run.return_value = None 222 mock_open_new_tab.return_value = None 223 mock_getsize.return_value = 0 224 mock_exists.return_value = True 225 terminal_output = io.StringIO() 226 sys.stdout = terminal_output 227 mock_process = mock_popen.return_value 228 mock_process.stdout = BytesIO(b'Trace loaded') 229 230 error = open_trace(TEST_FILE, WEB_UI_ADDRESS, True) 231 232 mock_open_new_tab.assert_called() 233 self.assertEqual(error, None) 234 self.assertEqual(terminal_output.getvalue(), 235 "\033[93m##### Loading trace. #####\n##### " 236 "Follow the directions in the Perfetto UI. Do not exit " 237 "out of torq until you are done viewing the trace. Press " 238 "CTRL+C to exit torq and close the trace_processor. " 239 "#####\033[0m\nProcess was killed.\n") 240 241 @mock.patch.dict(os.environ, {"ANDROID_BUILD_TOP": ANDROID_BUILD_TOP}, 242 clear=True) 243 @mock.patch.object(os.path, "exists", autospec=True) 244 @mock.patch.object(os.path, "getsize", autospec=True) 245 def test_download_trace_processor_android_scripts_exist(self, mock_getsize, 246 mock_exists): 247 mock_getsize.return_value = LARGE_FILE_SIZE 248 mock_exists.return_value = True 249 terminal_output = io.StringIO() 250 sys.stdout = terminal_output 251 252 trace_processor_path = download_trace_processor(TEST_FILE) 253 254 self.assertEqual(trace_processor_path, "%s%s" 255 % (ANDROID_BUILD_TOP, ANDROID_TRACE_PROCESSOR)) 256 self.assertEqual(terminal_output.getvalue(), "") 257 258 @mock.patch.dict(os.environ, {}, clear=True) 259 @mock.patch.object(os.path, "exists", autospec=True) 260 @mock.patch.object(os.path, "getsize", autospec=True) 261 def test_download_trace_processor_temp_scripts_exist(self, mock_getsize, 262 mock_exists): 263 mock_getsize.return_value = LARGE_FILE_SIZE 264 mock_exists.return_value = True 265 terminal_output = io.StringIO() 266 sys.stdout = terminal_output 267 268 trace_processor_path = download_trace_processor(TEST_FILE) 269 270 self.assertEqual(trace_processor_path, TORQ_TEMP_TRACE_PROCESSOR) 271 self.assertEqual(terminal_output.getvalue(), "") 272 273if __name__ == '__main__': 274 unittest.main() 275