[web] Add tests for StaticViewer and minor fixes.

mitmproxy/addons/static_viewer.py

@@ -42,24 +42,27 @@ def save_flows(path: pathlib.Path, flows: typing.Iterable[flow.Flow]) -> None:
 
 
 def save_flows_content(path: pathlib.Path, flows: typing.Iterable[flow.Flow]) -> None:
-    for flow in flows:
+    for f in flows:
         for m in ('request', 'response'):
-            message = getattr(flow, m)
-            message_path = path / "flows" / flow.id / m
+            message = getattr(f, m)
+            message_path = path / "flows" / f.id / m
             os.makedirs(str(message_path / "content"), exist_ok=True)
 
-            with open(str(message_path / '_content'), 'wb') as f:
+            with open(str(message_path / '_content'), 'wb') as content_file:
                 # don't use raw_content here as this is served with a default content type
-                if not message:
-                    # skip missing message
-                    continue
-                f.write(message.content)
+                if message:
+                    content_file.write(message.content)
+                else:
+                    content_file.write(b'No content.')
 
             # content_view
             t = time.time()
-            description, lines, error = contentviews.get_message_content_view(
-                'Auto', message
-            )
+            if message:
+                description, lines, error = contentviews.get_message_content_view(
+                    'Auto', message
+                )
+            else:
+                description, lines = 'No content.', []
             if time.time() - t > 0.1:
                 ctx.log(
                     "Slow content view: {} took {}s".format(
@@ -68,10 +71,10 @@ def save_flows_content(path: pathlib.Path, flows: typing.Iterable[flow.Flow]) ->
                     ),
                     "info"
                 )
-            with open(str(message_path / "content" / "Auto.json"), "w") as f:
+            with open(str(message_path / "content" / "Auto.json"), "w") as content_view_file:
                 json.dump(
                     dict(lines=list(lines), description=description),
-                    f
+                    content_view_file
                 )
 
 

test/mitmproxy/addons/test_static_viewer.py

@@ -0,0 +1,60 @@
+import json
+
+from mitmproxy.test import taddons
+from mitmproxy.test import tflow
+
+from mitmproxy import flowfilter
+from mitmproxy.tools.web.app import flow_to_json
+
+from mitmproxy.addons import static_viewer
+from mitmproxy.addons import save
+
+
+def test_save_static(tmpdir):
+    tmpdir.mkdir('static')
+    static_viewer.save_static(tmpdir)
+    assert len(tmpdir.listdir()) == 2
+    assert tmpdir.join('index.html').check(file=1)
+    assert tmpdir.join('static/static.js').read() == 'MITMWEB_STATIC = true;'
+
+
+def test_save_filter_help(tmpdir):
+    static_viewer.save_filter_help(tmpdir)
+    f = tmpdir.join('/filter-help.json')
+    assert f.check(file=1)
+    assert f.read() == json.dumps(dict(commands=flowfilter.help))
+
+
+def test_save_flows(tmpdir):
+    flows = [tflow.tflow(req=True, resp=None), tflow.tflow(req=True, resp=True)]
+    static_viewer.save_flows(tmpdir, flows)
+    assert tmpdir.join('flows.json').check(file=1)
+    assert tmpdir.join('flows.json').read() == json.dumps([flow_to_json(f) for f in flows])
+
+
+def test_save_flows_content(tmpdir):
+    flows = [tflow.tflow(req=True, resp=None), tflow.tflow(req=True, resp=True)]
+    static_viewer.save_flows_content(tmpdir, flows)
+    flows_path = tmpdir.join('flows')
+    assert len(flows_path.listdir()) == len(flows)
+    for p in flows_path.listdir():
+        assert p.join('request').check(dir=1)
+        assert p.join('response').check(dir=1)
+        assert p.join('request/_content').check(file=1)
+        assert p.join('request/content').check(dir=1)
+        assert p.join('response/_content').check(file=1)
+        assert p.join('response/content').check(dir=1)
+        assert p.join('request/content/Auto.json').check(file=1)
+        assert p.join('response/content/Auto.json').check(file=1)
+
+
+def test_static_viewer(tmpdir):
+    s = static_viewer.StaticViewer()
+    sa = save.Save()
+    with taddons.context() as tctx:
+        sa.save([tflow.tflow(resp=True)], str(tmpdir.join('foo')))
+        tctx.master.addons.add(s)
+        tctx.configure(s, web_static_viewer=str(tmpdir), rfile=str(tmpdir.join('foo')))
+        assert tmpdir.join('index.html').check(file=1)
+        assert tmpdir.join('static').check(dir=1)
+        assert tmpdir.join('flows').check(dir=1)