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.test import tflow
from mitmproxy.net.http import http1 from mitmproxy.net.http import http1
from mitmproxy.net.tcp import TCPClient 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 """head:'%s/p/200:h"Content-Length"="42"'""" % self.server.urlbase
) )
assert resp.headers["Content-Length"] == "42" 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)