1 files changed, 196 insertions, 0 deletions
diff --git a/testing/mozbase/mozproxy/tests/test_proxy.py b/testing/mozbase/mozproxy/tests/test_proxy.py
new file mode 100644
index 0000000000..47da3117ba
--- /dev/null
+++ b/testing/mozbase/mozproxy/tests/test_proxy.py
@@ -0,0 +1,196 @@
+#!/usr/bin/env python
+import os
+from unittest import mock
+
+import mozinfo
+import mozunit
+import requests
+from mozproxy import get_playback
+from support import tempdir
+
+here = os.path.dirname(__file__)
+
+
+class Process:
+    def __init__(self, *args, **kw):
+        pass
+
+    def run(self):
+        print("I am running something")
+
+    def poll(self):
+        return None
+
+    def wait(self):
+        return 0
+
+    def kill(self, sig=9):
+        pass
+
+    proc = object()
+    pid = 1234
+    stderr = stdout = None
+    returncode = 0
+
+
+_RETRY = 0
+
+
+class ProcessWithRetry(Process):
+    def __init__(self, *args, **kw):
+        Process.__init__(self, *args, **kw)
+
+    def wait(self):
+        global _RETRY
+        _RETRY += 1
+        if _RETRY >= 2:
+            _RETRY = 0
+            return 0
+        return -1
+
+
+def kill(pid, signal):
+    if pid == 1234:
+        return
+    return os.kill(pid, signal)
+
+
+def get_status_code(url, playback):
+    response = requests.get(
+        url=url, proxies={"http": "http://%s:%s/" % (playback.host, playback.port)}
+    )
+    return response.status_code
+
+
+def test_mitm_check_proxy(*args):
+    # test setup
+    pageset_name = os.path.join(here, "files", "mitm5-linux-firefox-amazon.manifest")
+
+    config = {
+        "playback_tool": "mitmproxy",
+        "playback_files": [os.path.join(here, "files", pageset_name)],
+        "playback_version": "8.1.1",
+        "platform": mozinfo.os,
+        "run_local": "MOZ_AUTOMATION" not in os.environ,
+        "binary": "firefox",
+        "app": "firefox",
+        "host": "127.0.0.1",
+    }
+
+    with tempdir() as obj_path:
+        config["obj_path"] = obj_path
+        playback = get_playback(config)
+        assert playback is not None
+
+        try:
+            playback.start()
+
+            url = "https://m.media-amazon.com/images/G/01/csm/showads.v2.js"
+            assert get_status_code(url, playback) == 200
+
+            url = "http://mozproxy/checkProxy"
+            assert get_status_code(url, playback) == 404
+        finally:
+            playback.stop()
+
+
+@mock.patch("mozproxy.backends.mitm.Mitmproxy.check_proxy")
+@mock.patch("mozproxy.backends.mitm.mitm.ProcessHandler", new=Process)
+@mock.patch("mozproxy.utils.ProcessHandler", new=Process)
+@mock.patch("os.kill", new=kill)
+def test_mitm(*args):
+    pageset_name = os.path.join(here, "files", "mitm5-linux-firefox-amazon.manifest")
+
+    config = {
+        "playback_tool": "mitmproxy",
+        "playback_files": [pageset_name],
+        "playback_version": "8.1.1",
+        "platform": mozinfo.os,
+        "run_local": True,
+        "binary": "firefox",
+        "app": "firefox",
+        "host": "example.com",
+    }
+
+    with tempdir() as obj_path:
+        config["obj_path"] = obj_path
+        playback = get_playback(config)
+    assert playback is not None
+    try:
+        playback.start()
+    finally:
+        playback.stop()
+
+
+@mock.patch("mozproxy.backends.mitm.Mitmproxy.check_proxy")
+@mock.patch("mozproxy.backends.mitm.mitm.ProcessHandler", new=Process)
+@mock.patch("mozproxy.utils.ProcessHandler", new=Process)
+@mock.patch("os.kill", new=kill)
+def test_playback_setup_failed(*args):
+    class SetupFailed(Exception):
+        pass
+
+    def setup(*args, **kw):
+        def _s(self):
+            raise SetupFailed("Failed")
+
+        return _s
+
+    pageset_name = os.path.join(here, "files", "mitm5-linux-firefox-amazon.manifest")
+
+    config = {
+        "playback_tool": "mitmproxy",
+        "playback_files": [pageset_name],
+        "playback_version": "4.0.4",
+        "platform": mozinfo.os,
+        "run_local": True,
+        "binary": "firefox",
+        "app": "firefox",
+        "host": "example.com",
+    }
+
+    prefix = "mozproxy.backends.mitm.desktop.MitmproxyDesktop."
+
+    with tempdir() as obj_path:
+        config["obj_path"] = obj_path
+        with mock.patch(prefix + "setup", new_callable=setup):
+            with mock.patch(prefix + "stop_mitmproxy_playback") as p:
+                try:
+                    pb = get_playback(config)
+                    pb.start()
+                except SetupFailed:
+                    assert p.call_count == 1
+                except Exception:
+                    raise
+
+
+@mock.patch("mozproxy.backends.mitm.Mitmproxy.check_proxy")
+@mock.patch("mozproxy.backends.mitm.mitm.ProcessHandler", new=ProcessWithRetry)
+@mock.patch("mozproxy.utils.ProcessHandler", new=ProcessWithRetry)
+@mock.patch("os.kill", new=kill)
+def test_mitm_with_retry(*args):
+    pageset_name = os.path.join(here, "files", "mitm5-linux-firefox-amazon.manifest")
+
+    config = {
+        "playback_tool": "mitmproxy",
+        "playback_files": [pageset_name],
+        "playback_version": "8.1.1",
+        "platform": mozinfo.os,
+        "run_local": True,
+        "binary": "firefox",
+        "app": "firefox",
+        "host": "example.com",
+    }
+
+    with tempdir() as obj_path:
+        config["obj_path"] = obj_path
+        playback = get_playback(config)
+    assert playback is not None
+    try:
+        playback.start()
+    finally:
+        playback.stop()
+
+
+if __name__ == "__main__":
+    mozunit.main(runwith="pytest")