Test for streaming

This commit is contained in:
Ujjwal Verma 2017-06-12 23:46:43 +05:30
parent 33252cb6bb
commit 6ca45856b4

View File

@ -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)