test_proxy.py (5450B)
1 #!/usr/bin/env python 2 import os 3 from unittest import mock 4 5 import mozinfo 6 import mozunit 7 import requests 8 from mozproxy import get_playback 9 from support import tempdir 10 11 here = os.path.dirname(__file__) 12 13 14 class Process: 15 def __init__(self, *args, **kw): 16 pass 17 18 def run(self): 19 print("I am running something") 20 21 def poll(self): 22 return None 23 24 def communicate(self): 25 return (["mock stdout"], ["mock stderr"]) 26 27 def wait(self): 28 return 0 29 30 def kill(self, sig=9): 31 pass 32 33 proc = object() 34 pid = 1234 35 stderr = stdout = [] 36 returncode = 0 37 38 39 _RETRY = 0 40 41 42 class ProcessWithRetry(Process): 43 def __init__(self, *args, **kw): 44 Process.__init__(self, *args, **kw) 45 46 def wait(self): 47 global _RETRY 48 _RETRY += 1 49 if _RETRY >= 2: 50 _RETRY = 0 51 return 0 52 return -1 53 54 55 def kill(pid, signal): 56 if pid == 1234: 57 return 58 return os.kill(pid, signal) 59 60 61 def get_status_code(url, playback): 62 response = requests.get( 63 url=url, proxies={"http": "http://%s:%s/" % (playback.host, playback.port)} 64 ) 65 return response.status_code 66 67 68 def cleanup(): 69 # some tests create this file as a side-effect 70 policies_file = os.path.join("distribution", "policies.json") 71 try: 72 if os.path.exists(policies_file): 73 os.remove(policies_file) 74 except PermissionError: 75 pass 76 77 78 def test_mitm_check_proxy(*args): 79 # test setup 80 pageset_name = os.path.join(here, "files", "mitm5-linux-firefox-amazon.manifest") 81 82 config = { 83 "playback_tool": "mitmproxy", 84 "playback_files": [os.path.join(here, "files", pageset_name)], 85 "playback_version": "8.1.1", 86 "platform": mozinfo.os, 87 "run_local": "MOZ_AUTOMATION" not in os.environ, 88 "binary": "firefox", 89 "app": "firefox", 90 "host": "127.0.0.1", 91 } 92 93 with tempdir() as obj_path: 94 config["obj_path"] = obj_path 95 playback = get_playback(config) 96 assert playback is not None 97 98 try: 99 playback.start() 100 101 url = "https://m.media-amazon.com/images/G/01/csm/showads.v2.js" 102 assert get_status_code(url, playback) == 200 103 104 url = "http://mozproxy/checkProxy" 105 assert get_status_code(url, playback) == 404 106 finally: 107 playback.stop() 108 cleanup() 109 110 111 @mock.patch("mozproxy.backends.mitm.Mitmproxy.check_proxy") 112 @mock.patch("mozproxy.backends.mitm.mitm.ProcessHandler", new=Process) 113 @mock.patch("mozproxy.utils.Popen", new=Process) 114 @mock.patch("os.kill", new=kill) 115 def test_mitm(*args): 116 pageset_name = os.path.join(here, "files", "mitm5-linux-firefox-amazon.manifest") 117 118 config = { 119 "playback_tool": "mitmproxy", 120 "playback_files": [pageset_name], 121 "playback_version": "8.1.1", 122 "platform": mozinfo.os, 123 "run_local": True, 124 "binary": "firefox", 125 "app": "firefox", 126 "host": "example.com", 127 } 128 129 with tempdir() as obj_path: 130 config["obj_path"] = obj_path 131 playback = get_playback(config) 132 assert playback is not None 133 try: 134 playback.start() 135 finally: 136 playback.stop() 137 cleanup() 138 139 140 @mock.patch("mozproxy.backends.mitm.Mitmproxy.check_proxy") 141 @mock.patch("mozproxy.backends.mitm.mitm.ProcessHandler", new=Process) 142 @mock.patch("mozproxy.utils.Popen", new=Process) 143 @mock.patch("os.kill", new=kill) 144 def test_playback_setup_failed(*args): 145 class SetupFailed(Exception): 146 pass 147 148 def setup(*args, **kw): 149 def _s(self): 150 raise SetupFailed("Failed") 151 152 return _s 153 154 pageset_name = os.path.join(here, "files", "mitm5-linux-firefox-amazon.manifest") 155 156 config = { 157 "playback_tool": "mitmproxy", 158 "playback_files": [pageset_name], 159 "playback_version": "4.0.4", 160 "platform": mozinfo.os, 161 "run_local": True, 162 "binary": "firefox", 163 "app": "firefox", 164 "host": "example.com", 165 } 166 167 prefix = "mozproxy.backends.mitm.desktop.MitmproxyDesktop." 168 169 with tempdir() as obj_path: 170 config["obj_path"] = obj_path 171 with mock.patch(prefix + "setup", new_callable=setup): 172 with mock.patch(prefix + "stop_mitmproxy_playback") as p: 173 try: 174 pb = get_playback(config) 175 pb.start() 176 except SetupFailed: 177 assert p.call_count == 1 178 except Exception: 179 raise 180 181 182 @mock.patch("mozproxy.backends.mitm.Mitmproxy.check_proxy") 183 @mock.patch("mozproxy.backends.mitm.mitm.ProcessHandler", new=ProcessWithRetry) 184 @mock.patch("mozproxy.utils.Popen", new=ProcessWithRetry) 185 @mock.patch("os.kill", new=kill) 186 def test_mitm_with_retry(*args): 187 pageset_name = os.path.join(here, "files", "mitm5-linux-firefox-amazon.manifest") 188 189 config = { 190 "playback_tool": "mitmproxy", 191 "playback_files": [pageset_name], 192 "playback_version": "8.1.1", 193 "platform": mozinfo.os, 194 "run_local": True, 195 "binary": "firefox", 196 "app": "firefox", 197 "host": "example.com", 198 } 199 200 with tempdir() as obj_path: 201 config["obj_path"] = obj_path 202 playback = get_playback(config) 203 assert playback is not None 204 try: 205 playback.start() 206 finally: 207 playback.stop() 208 cleanup() 209 210 211 if __name__ == "__main__": 212 mozunit.main(runwith="pytest")