1 files changed, 165 insertions, 0 deletions
diff --git a/test/examples/webscanner_helper/test_mapping.py b/test/examples/webscanner_helper/test_mapping.py
new file mode 100644
index 00000000..e4d519fc
--- /dev/null
+++ b/test/examples/webscanner_helper/test_mapping.py
@@ -0,0 +1,165 @@
+from typing import TextIO, Callable
+from unittest import mock
+from unittest.mock import MagicMock
+
+from mitmproxy.test import tflow
+from mitmproxy.test import tutils
+
+from examples.complex.webscanner_helper.mapping import MappingAddon, MappingAddonConfig
+
+
+class TestConfig:
+
+    def test_config(self):
+        assert MappingAddonConfig.HTML_PARSER == "html.parser"
+
+
+url = "http://10.10.10.10"
+new_content = "My Text"
+mapping_content = f'{{"{url}": {{"body": "{new_content}"}}}}'
+
+
+class TestMappingAddon:
+
+    def test_init(self, tmpdir):
+        tmpfile = tmpdir.join("tmpfile")
+        with open(tmpfile, "w") as tfile:
+            tfile.write(mapping_content)
+        mapping = MappingAddon(tmpfile)
+        assert "My Text" in str(mapping.mapping_templates._dump())
+
+    def test_load(self, tmpdir):
+        tmpfile = tmpdir.join("tmpfile")
+        with open(tmpfile, "w") as tfile:
+            tfile.write(mapping_content)
+        mapping = MappingAddon(tmpfile)
+        loader = MagicMock()
+
+        mapping.load(loader)
+        assert 'mapping_file' in str(loader.add_option.call_args_list)
+        assert 'map_persistent' in str(loader.add_option.call_args_list)
+
+    def test_configure(self, tmpdir):
+        tmpfile = tmpdir.join("tmpfile")
+        with open(tmpfile, "w") as tfile:
+            tfile.write(mapping_content)
+        mapping = MappingAddon(tmpfile)
+        new_filename = "My new filename"
+        updated = {str(mapping.OPT_MAPPING_FILE): new_filename, str(mapping.OPT_MAP_PERSISTENT): True}
+
+        open_mock = mock.mock_open(read_data="{}")
+        with mock.patch("builtins.open", open_mock):
+            mapping.configure(updated)
+        assert new_filename in str(open_mock.mock_calls)
+        assert mapping.filename == new_filename
+        assert mapping.persistent
+
+    def test_response_filtered(self, tmpdir):
+        tmpfile = tmpdir.join("tmpfile")
+        with open(tmpfile, "w") as tfile:
+            tfile.write(mapping_content)
+        mapping = MappingAddon(tmpfile)
+        f = tflow.tflow(resp=tutils.tresp())
+        test_content = b"Test"
+        f.response.content = test_content
+
+        mapping.response(f)
+        assert f.response.content == test_content
+
+    def test_response(self, tmpdir):
+        tmpfile = tmpdir.join("tmpfile")
+        with open(tmpfile, "w") as tfile:
+            tfile.write(mapping_content)
+        mapping = MappingAddon(tmpfile)
+        f = tflow.tflow(resp=tutils.tresp())
+        test_content = b"<body> Test </body>"
+        f.response.content = test_content
+        f.request.url = url
+
+        mapping.response(f)
+        assert f.response.content.decode("utf-8") == new_content
+
+    def test_response_content_type(self, tmpdir):
+        tmpfile = tmpdir.join("tmpfile")
+        with open(tmpfile, "w") as tfile:
+            tfile.write(mapping_content)
+        mapping = MappingAddon(tmpfile)
+        f = tflow.tflow(resp=tutils.tresp())
+        test_content = b"<body> Test </body>"
+        f.response.content = test_content
+        f.request.url = url
+        f.response.headers.add("content-type", "content-type")
+
+        mapping.response(f)
+        assert f.response.content == test_content
+
+    def test_response_not_existing(self, tmpdir):
+        tmpfile = tmpdir.join("tmpfile")
+        with open(tmpfile, "w") as tfile:
+            tfile.write(mapping_content)
+        mapping = MappingAddon(tmpfile)
+        f = tflow.tflow(resp=tutils.tresp())
+        test_content = b"<title> Test </title>"
+        f.response.content = test_content
+        f.request.url = url
+        mapping.response(f)
+        assert f.response.content == test_content
+
+    def test_persistance_false(self, tmpdir):
+        tmpfile = tmpdir.join("tmpfile")
+        with open(tmpfile, "w") as tfile:
+            tfile.write(mapping_content)
+        mapping = MappingAddon(tmpfile)
+
+        open_mock = mock.mock_open(read_data="{}")
+        with mock.patch("builtins.open", open_mock):
+            mapping.done()
+        assert len(open_mock.mock_calls) == 0
+
+    def test_persistance_true(self, tmpdir):
+        tmpfile = tmpdir.join("tmpfile")
+        with open(tmpfile, "w") as tfile:
+            tfile.write(mapping_content)
+        mapping = MappingAddon(tmpfile, persistent=True)
+
+        open_mock = mock.mock_open(read_data="{}")
+        with mock.patch("builtins.open", open_mock):
+            mapping.done()
+        with open(tmpfile, "r") as tfile:
+            results = tfile.read()
+        assert len(open_mock.mock_calls) != 0
+        assert results == mapping_content
+
+    def test_persistance_true_add_content(self, tmpdir):
+        tmpfile = tmpdir.join("tmpfile")
+        with open(tmpfile, "w") as tfile:
+            tfile.write(mapping_content)
+        mapping = MappingAddon(tmpfile, persistent=True)
+
+        f = tflow.tflow(resp=tutils.tresp())
+        test_content = b"<title> Test </title>"
+        f.response.content = test_content
+        f.request.url = url
+
+        mapping.response(f)
+        mapping.done()
+        with open(tmpfile, "r") as tfile:
+            results = tfile.read()
+        assert mapping_content in results
+
+    def mock_dump(self, f: TextIO, value_dumper: Callable):
+        assert value_dumper(None) == "None"
+        try:
+            value_dumper("Test")
+        except RuntimeError:
+            assert True
+        else:
+            assert False
+
+    def test_dump(selfself, tmpdir):
+        tmpfile = tmpdir.join("tmpfile")
+        with open(tmpfile, "w") as tfile:
+            tfile.write("{}")
+        mapping = MappingAddon(tmpfile, persistent=True)
+        with mock.patch('examples.complex.webscanner_helper.urldict.URLDict.dump', selfself.mock_dump):
+            mapping.done()