test_command_line.py (7684B)
1 #!/usr/bin/env python 2 import json 3 import os 4 import re 5 import signal 6 import subprocess 7 import sys 8 import threading 9 import time 10 11 import mozunit 12 from buildconfig import topobjdir, topsrcdir 13 14 here = os.path.dirname(__file__) 15 16 if os.name == "nt": 17 PROCESS_CREATION_FLAGS = subprocess.CREATE_NEW_PROCESS_GROUP 18 else: 19 PROCESS_CREATION_FLAGS = 0 20 21 22 # This is copied from <python/mozperftest/mozperftest/utils.py>. It's copied 23 # instead of imported since mozperfest is Python 3, and this file is 24 # (currently) Python 2. 25 def _install_package(virtualenv_manager, package): 26 from pip._internal.req.constructors import install_req_from_line 27 28 req = install_req_from_line(package) 29 req.check_if_exists(use_user_site=False) 30 # already installed, check if it's in our venv 31 if req.satisfied_by is not None: 32 venv_site_lib = os.path.abspath( 33 os.path.join(virtualenv_manager.bin_path, "..", "lib") 34 ) 35 site_packages = os.path.abspath(req.satisfied_by.location) 36 if site_packages.startswith(venv_site_lib): 37 # already installed in this venv, we can skip 38 return 39 40 subprocess.check_call([ 41 virtualenv_manager.python_path, 42 "-m", 43 "pip", 44 "install", 45 package, 46 ]) 47 48 49 def _kill_mozproxy(pid): 50 kill_signal = getattr(signal, "CTRL_BREAK_EVENT", signal.SIGINT) 51 os.kill(pid, kill_signal) 52 53 54 class OutputHandler: 55 def __init__(self): 56 self.port = None 57 self.port_event = threading.Event() 58 59 def __call__(self, line): 60 line = line.rstrip(b"\r\n") 61 if not line.strip(): 62 return 63 line = line.decode("utf-8", errors="replace") 64 # Print the output we received so we have useful logs if a test fails. 65 print(line) 66 67 try: 68 data = json.loads(line) 69 except ValueError: 70 return 71 72 if isinstance(data, dict) and "action" in data: 73 # Retrieve the port number for the proxy server from the logs of 74 # our subprocess. 75 m = re.match(r"Proxy running on port (\d+)", data.get("message", "")) 76 if m: 77 self.port = m.group(1) 78 self.port_event.set() 79 80 def finished(self): 81 self.port_event.set() 82 83 84 def test_help(): 85 p = subprocess.run([sys.executable, "-m", "mozproxy", "--help"], check=False) 86 assert p.returncode == 0 87 88 89 def test_run_record_no_files(): 90 output_handler = OutputHandler() 91 p = subprocess.Popen( 92 [ 93 sys.executable, 94 "-m", 95 "mozproxy", 96 "--local", 97 "--mode=record", 98 "--binary=firefox", 99 "--topsrcdir=" + topsrcdir, 100 "--objdir=" + topobjdir, 101 ], 102 creationflags=PROCESS_CREATION_FLAGS, 103 stdout=subprocess.PIPE, 104 stderr=subprocess.STDOUT, 105 text=False, 106 ) 107 108 for line in p.stdout: 109 output_handler(line) 110 if output_handler.port_event.is_set(): 111 break 112 output_handler.finished() 113 # The first time we run mozproxy, we need to fetch mitmproxy, which can 114 # take a while... 115 assert output_handler.port_event.wait(120) is True 116 # Give mitmproxy a bit of time to start up so we can verify that it's 117 # actually running before we kill mozproxy. 118 time.sleep(5) 119 _kill_mozproxy(p.pid) 120 121 # Assert process raises error 122 assert p.wait(10) == 2 123 assert output_handler.port is None 124 125 126 def test_run_record_multiple_files(): 127 output_handler = OutputHandler() 128 p = subprocess.Popen( 129 [ 130 sys.executable, 131 "-m", 132 "mozproxy", 133 "--local", 134 "--mode=record", 135 "--binary=firefox", 136 "--topsrcdir=" + topsrcdir, 137 "--objdir=" + topobjdir, 138 os.path.join(here, "files", "new_record.zip"), 139 os.path.join(here, "files", "new_record2.zip"), 140 ], 141 creationflags=PROCESS_CREATION_FLAGS, 142 stdout=subprocess.PIPE, 143 stderr=subprocess.STDOUT, 144 text=False, 145 ) 146 147 for line in p.stdout: 148 output_handler(line) 149 if output_handler.port_event.is_set(): 150 break 151 output_handler.finished() 152 # The first time we run mozproxy, we need to fetch mitmproxy, which can 153 # take a while... 154 assert output_handler.port_event.wait(120) is True 155 # Give mitmproxy a bit of time to start up so we can verify that it's 156 # actually running before we kill mozproxy. 157 time.sleep(5) 158 _kill_mozproxy(p.pid) 159 160 assert p.wait(10) == 4 161 assert output_handler.port is None 162 163 164 def test_run_record(): 165 output_handler = OutputHandler() 166 p = subprocess.Popen( 167 [ 168 sys.executable, 169 "-m", 170 "mozproxy", 171 "--local", 172 "--mode=record", 173 "--binary=firefox", 174 "--topsrcdir=" + topsrcdir, 175 "--objdir=" + topobjdir, 176 os.path.join(here, "files", "record.zip"), 177 ], 178 creationflags=PROCESS_CREATION_FLAGS, 179 stdout=subprocess.PIPE, 180 stderr=subprocess.STDOUT, 181 text=False, 182 ) 183 184 for line in p.stdout: 185 output_handler(line) 186 if output_handler.port_event.is_set(): 187 break 188 output_handler.finished() 189 try: 190 # The first time we run mozproxy, we need to fetch mitmproxy, which can 191 # take a while... 192 assert output_handler.port_event.wait(120) is True 193 # Give mitmproxy a bit of time to start up so we can verify that it's 194 # actually running before we kill mozproxy. 195 time.sleep(5) 196 _kill_mozproxy(p.pid) 197 198 assert p.wait(10) == 0 199 assert output_handler.port is not None 200 finally: 201 os.remove(os.path.join(here, "files", "record.zip")) 202 203 204 def test_run_playback(): 205 output_handler = OutputHandler() 206 p = subprocess.Popen( 207 [ 208 sys.executable, 209 "-m", 210 "mozproxy", 211 "--local", 212 "--binary=firefox", 213 "--topsrcdir=" + topsrcdir, 214 "--objdir=" + topobjdir, 215 os.path.join(here, "files", "mitm5-linux-firefox-amazon.zip"), 216 ], 217 creationflags=PROCESS_CREATION_FLAGS, 218 stdout=subprocess.PIPE, 219 stderr=subprocess.STDOUT, 220 text=False, 221 ) 222 223 for line in p.stdout: 224 output_handler(line) 225 if output_handler.port_event.is_set(): 226 break 227 output_handler.finished() 228 # The first time we run mozproxy, we need to fetch mitmproxy, which can 229 # take a while... 230 assert output_handler.port_event.wait(120) is True 231 # Give mitmproxy a bit of time to start up so we can verify that it's 232 # actually running before we kill mozproxy. 233 time.sleep(5) 234 _kill_mozproxy(p.pid) 235 236 assert p.wait(10) == 0 237 assert output_handler.port is not None 238 239 240 def test_failure(): 241 output_handler = OutputHandler() 242 p = subprocess.Popen( 243 [ 244 sys.executable, 245 "-m", 246 "mozproxy", 247 "--local", 248 # Exclude some options here to trigger a command-line error. 249 os.path.join(here, "files", "mitm5-linux-firefox-amazon.zip"), 250 ], 251 creationflags=PROCESS_CREATION_FLAGS, 252 stdout=subprocess.PIPE, 253 stderr=subprocess.STDOUT, 254 text=False, 255 ) 256 for line in p.stdout: 257 output_handler(line) 258 if output_handler.port_event.is_set(): 259 break 260 output_handler.finished() 261 assert output_handler.port_event.wait(10) is True 262 assert p.wait(10) == 2 263 assert output_handler.port is None 264 265 266 if __name__ == "__main__": 267 mozunit.main(runwith="pytest")