Merge pull request #2393 from ujjwal96/streaming-test

Test for streaming

1 files changed, 37 insertions, 0 deletions

diff --git a/test/mitmproxy/proxy/protocol/test_http1.py b/test/mitmproxy/proxy/protocol/test_http1.py
index b642afb3..1eff8666 100644
--- a/test/mitmproxy/proxy/protocol/test_http1.py
+++ b/test/mitmproxy/proxy/protocol/test_http1.py
@@ -1,3 +1,7 @@
+from unittest import mock
+import pytest
+
+from mitmproxy import exceptions
 from mitmproxy.test import tflow
 from mitmproxy.net.http import http1
 from mitmproxy.net.tcp import TCPClient
@@ -77,3 +81,36 @@ class TestHeadContentLength(tservers.HTTPProxyTest):
                 """head:'%s/p/200:h"Content-Length"="42"'""" % self.server.urlbase
             )
         assert resp.headers["Content-Length"] == "42"
+
+
+class TestStreaming(tservers.HTTPProxyTest):
+
+    @pytest.mark.parametrize('streaming', [True, False])
+    def test_streaming(self, streaming):
+
+        class Stream:
+            def requestheaders(self, f):
+                f.request.stream = streaming
+
+            def responseheaders(self, f):
+                f.response.stream = streaming
+
+        def assert_write(self, v):
+            if streaming:
+                assert len(v) <= 4096
+            return self.o.write(v)
+
+        self.master.addons.add(Stream())
+        p = self.pathoc()
+        with p.connect():
+            with mock.patch("mitmproxy.net.tcp.Writer.write", side_effect=assert_write, autospec=True):
+                # response with 10000 bytes
+                r = p.request("post:'%s/p/200:b@10000'" % self.server.urlbase)
+                assert len(r.content) == 10000
+
+                if streaming:
+                    with pytest.raises(exceptions.HttpReadDisconnect):  # as the assertion in assert_write fails
+                        # request with 10000 bytes
+                        p.request("post:'%s/p/200':b@10000" % self.server.urlbase)
+                else:
+                    assert p.request("post:'%s/p/200':b@10000" % self.server.urlbase)