Merge pull request #833 from zbuc/contentview_scripts

Contentview scripts

3 files changed, 68 insertions, 0 deletions

diff --git a/test/test_contentview.py b/test/test_contentview.py
index 14adcd83..2a70b414 100644
--- a/test/test_contentview.py
+++ b/test/test_contentview.py
@@ -210,6 +210,21 @@ Larry
         assert "decoded gzip" in r[0]
         assert "Raw" in r[0]
 
+    def test_add_cv(self):
+        class TestContentView(cv.View):
+            name = "test"
+            prompt = ("t", "test")
+
+        tcv = TestContentView()
+        cv.add(tcv)
+
+        # repeated addition causes exception
+        tutils.raises(
+            ContentViewException,
+            cv.add,
+            tcv
+        )
+
 
 if pyamf:
     def test_view_amf_request():
diff --git a/test/test_custom_contentview.py b/test/test_custom_contentview.py
new file mode 100644
index 00000000..4b5a3e53
--- /dev/null
+++ b/test/test_custom_contentview.py
@@ -0,0 +1,52 @@
+from libmproxy import script, flow
+import libmproxy.contentviews as cv
+from netlib.http import Headers
+
+
+def test_custom_views():
+    class ViewNoop(cv.View):
+        name = "noop"
+        prompt = ("noop", "n")
+        content_types = ["text/none"]
+
+        def __call__(self, data, **metadata):
+            return "noop", cv.format_text(data)
+
+
+    view_obj = ViewNoop()
+
+    cv.add(view_obj)
+
+    assert cv.get("noop")
+
+    r = cv.get_content_view(
+        cv.get("noop"),
+        "[1, 2, 3]",
+        headers=Headers(
+            content_type="text/plain"
+        )
+    )
+    assert "noop" in r[0]
+
+    # now try content-type matching
+    r = cv.get_content_view(
+        cv.get("Auto"),
+        "[1, 2, 3]",
+        headers=Headers(
+            content_type="text/none"
+        )
+    )
+    assert "noop" in r[0]
+
+    # now try removing the custom view
+    cv.remove(view_obj)
+    r = cv.get_content_view(
+        cv.get("Auto"),
+        "[1, 2, 3]",
+        headers=Headers(
+            content_type="text/none"
+        )
+    )
+    assert "noop" not in r[0]
+
+    
diff --git a/test/test_script.py b/test/test_script.py
index 1b0e5a5b..8612d5f3 100644
--- a/test/test_script.py
+++ b/test/test_script.py
@@ -127,3 +127,4 @@ def test_command_parsing():
     absfilepath = os.path.normcase(tutils.test_data.path("scripts/a.py"))
     s = script.Script(absfilepath, fm)
     assert os.path.isfile(s.args[0])
+