Response body streaming to addon

2 files changed, 29 insertions, 12 deletions

diff --git a/test/mitmproxy/builtins/test_streambodies.py b/test/mitmproxy/builtins/test_streambodies.py
new file mode 100644
index 00000000..0e8a82f2
--- /dev/null
+++ b/test/mitmproxy/builtins/test_streambodies.py
@@ -0,0 +1,29 @@
+from .. import tutils, mastertest
+from mitmproxy.flow import state
+from mitmproxy.flow import master
+from mitmproxy import options
+
+from mitmproxy.builtins import streambodies
+
+
+class TestStreamBodies(mastertest.MasterTest):
+    def test_simple(self):
+        s = state.DummyState()
+        o = options.Options(stream_large_bodies = 10)
+        m = master.FlowMaster(o, None, s)
+        sa = streambodies.StreamBodies()
+        m.addons.add(sa)
+
+        f = tutils.tflow()
+        f.request.content = b""
+        f.request.headers["Content-Length"] = "1024"
+        assert not f.request.stream
+        m.requestheaders(f)
+        assert f.request.stream
+
+        f = tutils.tflow(resp=True)
+        f.response.content = b""
+        f.response.headers["Content-Length"] = "1024"
+        assert not f.response.stream
+        m.responseheaders(f)
+        assert f.response.stream
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index a175d193..e9db210e 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -288,18 +288,6 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin, AppMixin):
             resp = p.request("get:'http://foo':h':foo'='bar'")
         assert resp.status_code == 400
 
-    def test_stream(self):
-        self.master.set_stream_large_bodies(1024 * 2)
-
-        self.pathod("200:b@1k")
-        assert not self.master.state.view[-1].response.stream
-        assert len(self.master.state.view[-1].response.content) == 1024 * 1
-
-        self.pathod("200:b@3k")
-        assert self.master.state.view[-1].response.stream
-        assert self.master.state.view[-1].response.content is None
-        self.master.set_stream_large_bodies(None)
-
     def test_stream_modify(self):
         s = script.Script(
             tutils.test_data.path("data/addonscripts/stream_modify.py")