Merge pull request #1706 from cortesi/taddons

addons tests and bugfixes

4 files changed, 69 insertions, 72 deletions

diff --git a/test/mitmproxy/addons/test_filestreamer.py b/test/mitmproxy/addons/test_filestreamer.py
deleted file mode 100644
index 28094c43..00000000
--- a/test/mitmproxy/addons/test_filestreamer.py
+++ /dev/null
@@ -1,44 +0,0 @@
-from mitmproxy.test import tflow
-from mitmproxy.test import tutils
-
-from .. import mastertest
-
-import os.path
-
-from mitmproxy.addons import filestreamer
-from mitmproxy import master
-from mitmproxy import io
-from mitmproxy import options
-from mitmproxy import proxy
-
-
-class TestStream(mastertest.MasterTest):
-    def test_stream(self):
-        with tutils.tmpdir() as tdir:
-            p = os.path.join(tdir, "foo")
-
-            def r():
-                r = io.FlowReader(open(p, "rb"))
-                return list(r.stream())
-
-            o = options.Options(
-                outfile = (p, "wb")
-            )
-            m = master.Master(o, proxy.DummyServer())
-            sa = filestreamer.FileStreamer()
-
-            m.addons.add(sa)
-            f = tflow.tflow(resp=True)
-            m.request(f)
-            m.response(f)
-            m.addons.remove(sa)
-
-            assert r()[0].response
-
-            m.options.outfile = (p, "ab")
-
-            m.addons.add(sa)
-            f = tflow.tflow()
-            m.request(f)
-            m.addons.remove(sa)
-            assert not r()[1].response
diff --git a/test/mitmproxy/addons/test_intercept.py b/test/mitmproxy/addons/test_intercept.py
index efdac374..a347f9ab 100644
--- a/test/mitmproxy/addons/test_intercept.py
+++ b/test/mitmproxy/addons/test_intercept.py
@@ -36,3 +36,8 @@ def test_simple():
         f = tflow.tflow(resp=False)
         tctx.cycle(r, f)
         assert not f.intercepted
+
+        f = tflow.tflow(resp=True)
+        f.reply._state = "handled"
+        r.response(f)
+        assert f.intercepted
diff --git a/test/mitmproxy/addons/test_streamfile.py b/test/mitmproxy/addons/test_streamfile.py
new file mode 100644
index 00000000..82a4345b
--- /dev/null
+++ b/test/mitmproxy/addons/test_streamfile.py
@@ -0,0 +1,64 @@
+from mitmproxy.test import tflow
+from mitmproxy.test import tutils
+from mitmproxy.test import taddons
+
+import os.path
+from mitmproxy import io
+from mitmproxy import exceptions
+from mitmproxy.tools import dump
+from mitmproxy.addons import streamfile
+
+
+def test_configure():
+    sa = streamfile.StreamFile()
+    with taddons.context(options=dump.Options()) as tctx:
+        with tutils.tmpdir() as tdir:
+            p = os.path.join(tdir, "foo")
+            tutils.raises(
+                exceptions.OptionsError,
+                tctx.configure, sa, streamfile=tdir
+            )
+            tutils.raises(
+                "invalid filter",
+                tctx.configure, sa, streamfile=p, filtstr="~~"
+            )
+
+
+def rd(p):
+    x = io.FlowReader(open(p, "rb"))
+    return list(x.stream())
+
+
+def test_tcp():
+    sa = streamfile.StreamFile()
+    with taddons.context() as tctx:
+        with tutils.tmpdir() as tdir:
+            p = os.path.join(tdir, "foo")
+            tctx.configure(sa, streamfile=p)
+
+            tt = tflow.ttcpflow()
+            sa.tcp_start(tt)
+            sa.tcp_end(tt)
+            tctx.configure(sa, streamfile=None)
+            assert rd(p)
+
+
+def test_simple():
+    sa = streamfile.StreamFile()
+    with taddons.context() as tctx:
+        with tutils.tmpdir() as tdir:
+            p = os.path.join(tdir, "foo")
+
+            tctx.configure(sa, streamfile=p)
+
+            f = tflow.tflow(resp=True)
+            sa.request(f)
+            sa.response(f)
+            tctx.configure(sa, streamfile=None)
+            assert rd(p)[0].response
+
+            tctx.configure(sa, streamfile=p, streamfile_append=True)
+            f = tflow.tflow()
+            sa.request(f)
+            tctx.configure(sa, streamfile=None)
+            assert not rd(p)[1].response
diff --git a/test/mitmproxy/test_dump.py b/test/mitmproxy/test_dump.py
index aa3228e4..e331637d 100644
--- a/test/mitmproxy/test_dump.py
+++ b/test/mitmproxy/test_dump.py
@@ -2,7 +2,6 @@ from mitmproxy.test import tflow
 import os
 import io
 
-import mitmproxy.io
 from mitmproxy.tools import dump
 from mitmproxy import exceptions
 from mitmproxy import proxy
@@ -126,33 +125,6 @@ class TestDumpMaster(mastertest.MasterTest):
         f = self.cycle(m, b"content")
         assert f.request.headers["one"] == "two"
 
-    def test_write(self):
-        with tutils.tmpdir() as d:
-            p = os.path.join(d, "a")
-            self.dummy_cycle(
-                self.mkmaster(None, outfile=(p, "wb"), verbosity=0), 1, b""
-            )
-            assert len(list(mitmproxy.io.FlowReader(open(p, "rb")).stream())) == 1
-
-    def test_write_append(self):
-        with tutils.tmpdir() as d:
-            p = os.path.join(d, "a.append")
-            self.dummy_cycle(
-                self.mkmaster(None, outfile=(p, "wb"), verbosity=0),
-                1, b""
-            )
-            self.dummy_cycle(
-                self.mkmaster(None, outfile=(p, "ab"), verbosity=0),
-                1, b""
-            )
-            assert len(list(mitmproxy.io.FlowReader(open(p, "rb")).stream())) == 2
-
-    def test_write_err(self):
-        tutils.raises(
-            exceptions.OptionsError,
-            self.mkmaster, None, outfile = ("nonexistentdir/foo", "wb")
-        )
-
     def test_script(self):
         ret = self.dummy_cycle(
             self.mkmaster(