1 files changed, 89 insertions, 0 deletions
diff --git a/testing/mozbase/mozproxy/tests/test_mitm_addons.py b/testing/mozbase/mozproxy/tests/test_mitm_addons.py
new file mode 100644
index 0000000000..ed3805ef9d
--- /dev/null
+++ b/testing/mozbase/mozproxy/tests/test_mitm_addons.py
@@ -0,0 +1,89 @@
+#!/usr/bin/env python
+import json
+import os
+from unittest import mock
+
+import mozunit
+
+here = os.path.dirname(__file__)
+os.environ["MOZPROXY_DIR"] = os.path.join(here, "files")
+
+protocol = {
+    "http_protocol": {"aax-us-iad.amazon.com": "HTTP/1.1"},
+    "recorded_requests": 4,
+    "recorded_requests_unique": 1,
+}
+
+
+@mock.patch(
+    "mozproxy.backends.mitm.scripts.http_protocol_extractor.HttpProtocolExtractor.get_ctx"
+)
+def test_http_protocol_generate_json_file(ctx_mock):
+    ctx_mock.return_value.options.save_stream_file = os.path.join(
+        os.environ["MOZPROXY_DIR"], "http_protocol_recording_done.mp"
+    )
+
+    from mozproxy.backends.mitm.scripts.http_protocol_extractor import (
+        HttpProtocolExtractor,
+    )
+
+    test_http_protocol = HttpProtocolExtractor()
+    test_http_protocol.ctx = test_http_protocol.get_ctx()
+
+    # test data
+    test_http_protocol.request_protocol = protocol["http_protocol"]
+    test_http_protocol.hashes = ["Hash string"]
+    test_http_protocol.request_count = protocol["recorded_requests"]
+
+    test_http_protocol.done()
+
+    json_path = os.path.join(
+        os.environ["MOZPROXY_DIR"], "http_protocol_recording_done.json"
+    )
+    assert os.path.exists(json_path)
+    with open(json_path) as json_file:
+        output_data = json.load(json_file)
+
+        assert output_data["recorded_requests"] == protocol["recorded_requests"]
+        assert (
+            output_data["recorded_requests_unique"]
+            == protocol["recorded_requests_unique"]
+        )
+        assert output_data["http_protocol"] == protocol["http_protocol"]
+
+
+@mock.patch(
+    "mozproxy.backends.mitm.scripts.http_protocol_extractor.HttpProtocolExtractor.get_ctx"
+)
+def test_http_protocol_response(ctx_mock):
+    ctx_mock.return_value.options.save_stream_file = os.path.join(
+        os.environ["MOZPROXY_DIR"], "http_protocol_recording_done.mp"
+    )
+
+    from mozproxy.backends.mitm.scripts.http_protocol_extractor import (
+        HttpProtocolExtractor,
+    )
+
+    test_http_protocol = HttpProtocolExtractor()
+    test_http_protocol.ctx = test_http_protocol.get_ctx()
+
+    # test data
+    flow = mock.MagicMock()
+    flow.type = "http"
+    flow.request.url = "https://www.google.com/complete/search"
+    flow.request.port = 33
+    flow.response.data.http_version = b"HTTP/1.1"
+
+    test_http_protocol.request_protocol = {}
+    test_http_protocol.hashes = []
+    test_http_protocol.request_count = 0
+
+    test_http_protocol.response(flow)
+
+    assert test_http_protocol.request_count == 1
+    assert len(test_http_protocol.hashes) == 1
+    assert test_http_protocol.request_protocol["www.google.com"] == "HTTP/1.1"
+
+
+if __name__ == "__main__":
+    mozunit.main(runwith="pytest")