diff options
author | Aldo Cortesi <aldo@nullcube.com> | 2016-07-14 12:43:17 +1200 |
---|---|---|
committer | Aldo Cortesi <aldo@nullcube.com> | 2016-07-14 13:52:09 +1200 |
commit | b2c4f301cb834ecdf6e5b0063e86be877d3ece6d (patch) | |
tree | 3496f88cbcfa1e7c2470e3e1bed9ac67bb563bdb /test | |
parent | 703c05066ec0bc05c680e24d368606791dd1c958 (diff) | |
download | mitmproxy-b2c4f301cb834ecdf6e5b0063e86be877d3ece6d.tar.gz mitmproxy-b2c4f301cb834ecdf6e5b0063e86be877d3ece6d.tar.bz2 mitmproxy-b2c4f301cb834ecdf6e5b0063e86be877d3ece6d.zip |
Stream to file -> addon
This commit also clarifies a confusion about the "outfile" attribute and its
use in testing in the mitmdump master.
Diffstat (limited to 'test')
-rw-r--r-- | test/mitmproxy/builtins/test_stream.py | 46 | ||||
-rw-r--r-- | test/mitmproxy/mastertest.py | 13 | ||||
-rw-r--r-- | test/mitmproxy/test_dump.py | 89 | ||||
-rw-r--r-- | test/mitmproxy/test_flow.py | 31 |
4 files changed, 104 insertions, 75 deletions
diff --git a/test/mitmproxy/builtins/test_stream.py b/test/mitmproxy/builtins/test_stream.py new file mode 100644 index 00000000..54e4f7d9 --- /dev/null +++ b/test/mitmproxy/builtins/test_stream.py @@ -0,0 +1,46 @@ +from __future__ import absolute_import, print_function, division + +from .. import tutils, mastertest + +import os.path + +from mitmproxy.builtins import stream +from mitmproxy.flow import master, FlowReader +from mitmproxy.flow import state +from mitmproxy import options + + +class TestStream(mastertest.MasterTest): + def test_stream(self): + with tutils.tmpdir() as tdir: + p = os.path.join(tdir, "foo") + + def r(): + r = FlowReader(open(p, "rb")) + return list(r.stream()) + + s = state.State() + m = master.FlowMaster( + options.Options( + outfile = (p, "wb") + ), + None, + s + ) + sa = stream.Stream() + + m.addons.add(sa) + f = tutils.tflow(resp=True) + self.invoke(m, "request", f) + self.invoke(m, "response", f) + m.addons.remove(sa) + + assert r()[0].response + + m.options.outfile = (p, "ab") + + m.addons.add(sa) + f = tutils.tflow() + self.invoke(m, "request", f) + m.addons.remove(sa) + assert not r()[1].response diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py index 06854e25..9754d3a9 100644 --- a/test/mitmproxy/mastertest.py +++ b/test/mitmproxy/mastertest.py @@ -18,15 +18,14 @@ class MasterTest: l = proxy.Log("connect") l.reply = mock.MagicMock() master.log(l) - master.clientconnect(f.client_conn) - master.serverconnect(f.server_conn) - master.request(f) + self.invoke(master, "clientconnect", f.client_conn) + self.invoke(master, "clientconnect", f.client_conn) + self.invoke(master, "serverconnect", f.server_conn) + self.invoke(master, "request", f) if not f.error: f.response = models.HTTPResponse.wrap(netlib.tutils.tresp(content=content)) - f.reply.acked = False - f = master.response(f) - f.client_conn.reply.acked = False - master.clientdisconnect(f.client_conn) + self.invoke(master, "response", f) + self.invoke(master, "clientdisconnect", f) return f def dummy_cycle(self, master, n, content): diff --git a/test/mitmproxy/test_dump.py b/test/mitmproxy/test_dump.py index aa73b5a4..9686be84 100644 --- a/test/mitmproxy/test_dump.py +++ b/test/mitmproxy/test_dump.py @@ -4,31 +4,33 @@ from mitmproxy.exceptions import ContentViewException import netlib.tutils -from mitmproxy import dump, flow, models +from mitmproxy import dump, flow, models, exceptions from . import tutils, mastertest import mock def test_strfuncs(): - o = dump.Options() + o = dump.Options( + tfile = StringIO(), + flow_detail = 0, + ) m = dump.DumpMaster(None, o) - m.outfile = StringIO() m.o.flow_detail = 0 m.echo_flow(tutils.tflow()) - assert not m.outfile.getvalue() + assert not o.tfile.getvalue() m.o.flow_detail = 4 m.echo_flow(tutils.tflow()) - assert m.outfile.getvalue() + assert o.tfile.getvalue() - m.outfile = StringIO() + o.tfile = StringIO() m.echo_flow(tutils.tflow(resp=True)) - assert "<<" in m.outfile.getvalue() + assert "<<" in o.tfile.getvalue() - m.outfile = StringIO() + o.tfile = StringIO() m.echo_flow(tutils.tflow(err=True)) - assert "<<" in m.outfile.getvalue() + assert "<<" in o.tfile.getvalue() flow = tutils.tflow() flow.request = netlib.tutils.treq() @@ -50,25 +52,32 @@ def test_strfuncs(): def test_contentview(get_content_view): get_content_view.side_effect = ContentViewException(""), ("x", iter([])) - o = dump.Options(flow_detail=4, verbosity=3) - m = dump.DumpMaster(None, o, StringIO()) + o = dump.Options( + flow_detail=4, + verbosity=3, + tfile=StringIO(), + ) + m = dump.DumpMaster(None, o) m.echo_flow(tutils.tflow()) - assert "Content viewer failed" in m.outfile.getvalue() + assert "Content viewer failed" in m.options.tfile.getvalue() class TestDumpMaster(mastertest.MasterTest): def dummy_cycle(self, master, n, content): mastertest.MasterTest.dummy_cycle(self, master, n, content) - return master.outfile.getvalue() + return master.options.tfile.getvalue() def mkmaster(self, filt, **options): - cs = StringIO() if "verbosity" not in options: options["verbosity"] = 0 if "flow_detail" not in options: options["flow_detail"] = 0 - o = dump.Options(filtstr=filt, **options) - return dump.DumpMaster(None, o, outfile=cs) + o = dump.Options( + filtstr=filt, + tfile=StringIO(), + **options + ) + return dump.DumpMaster(None, o) def test_basic(self): for i in (1, 2, 3): @@ -89,31 +98,33 @@ class TestDumpMaster(mastertest.MasterTest): ) def test_error(self): - cs = StringIO() - o = dump.Options(flow_detail=1) - m = dump.DumpMaster(None, o, outfile=cs) + o = dump.Options( + tfile=StringIO(), + flow_detail=1 + ) + m = dump.DumpMaster(None, o) f = tutils.tflow(err=True) m.request(f) assert m.error(f) - assert "error" in cs.getvalue() + assert "error" in o.tfile.getvalue() def test_missing_content(self): - cs = StringIO() - o = dump.Options(flow_detail=3) - m = dump.DumpMaster(None, o, outfile=cs) + o = dump.Options( + flow_detail=3, + tfile=StringIO(), + ) + m = dump.DumpMaster(None, o) f = tutils.tflow() f.request.content = None m.request(f) f.response = models.HTTPResponse.wrap(netlib.tutils.tresp()) f.response.content = None m.response(f) - assert "content missing" in cs.getvalue() + assert "content missing" in o.tfile.getvalue() def test_replay(self): - cs = StringIO() - o = dump.Options(server_replay=["nonexistent"], kill=True) - tutils.raises(dump.DumpError, dump.DumpMaster, None, o, outfile=cs) + tutils.raises(dump.DumpError, dump.DumpMaster, None, o) with tutils.tmpdir() as t: p = os.path.join(t, "rep") @@ -122,7 +133,7 @@ class TestDumpMaster(mastertest.MasterTest): o = dump.Options(server_replay=[p], kill=True) o.verbosity = 0 o.flow_detail = 0 - m = dump.DumpMaster(None, o, outfile=cs) + m = dump.DumpMaster(None, o) self.cycle(m, b"content") self.cycle(m, b"content") @@ -130,13 +141,13 @@ class TestDumpMaster(mastertest.MasterTest): o = dump.Options(server_replay=[p], kill=False) o.verbosity = 0 o.flow_detail = 0 - m = dump.DumpMaster(None, o, outfile=cs) + m = dump.DumpMaster(None, o) self.cycle(m, b"nonexistent") o = dump.Options(client_replay=[p], kill=False) o.verbosity = 0 o.flow_detail = 0 - m = dump.DumpMaster(None, o, outfile=cs) + m = dump.DumpMaster(None, o) def test_read(self): with tutils.tmpdir() as t: @@ -172,20 +183,24 @@ class TestDumpMaster(mastertest.MasterTest): assert len(m.apps.apps) == 1 def test_replacements(self): - cs = StringIO() - o = dump.Options(replacements=[(".*", "content", "foo")]) + o = dump.Options( + replacements=[(".*", "content", "foo")], + tfile = StringIO(), + ) o.verbosity = 0 o.flow_detail = 0 - m = dump.DumpMaster(None, o, outfile=cs) + m = dump.DumpMaster(None, o) f = self.cycle(m, b"content") assert f.request.content == b"foo" def test_setheader(self): - cs = StringIO() - o = dump.Options(setheaders=[(".*", "one", "two")]) + o = dump.Options( + setheaders=[(".*", "one", "two")], + tfile=StringIO() + ) o.verbosity = 0 o.flow_detail = 0 - m = dump.DumpMaster(None, o, outfile=cs) + m = dump.DumpMaster(None, o) f = self.cycle(m, b"content") assert f.request.headers["one"] == "two" @@ -212,7 +227,7 @@ class TestDumpMaster(mastertest.MasterTest): def test_write_err(self): tutils.raises( - dump.DumpError, + exceptions.OptionsError, self.mkmaster, None, outfile = ("nonexistentdir/foo", "wb") ) diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 91342e58..1a07f74d 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -1,5 +1,3 @@ -import os.path - import mock import io @@ -887,35 +885,6 @@ class TestFlowMaster: fm.process_new_request(f) assert "killed" in f.error.msg - def test_stream(self): - with tutils.tmpdir() as tdir: - p = os.path.join(tdir, "foo") - - def read(): - with open(p, "rb") as f: - r = flow.FlowReader(f) - return list(r.stream()) - - s = flow.State() - fm = flow.FlowMaster(None, None, s) - f = tutils.tflow(resp=True) - - with open(p, "ab") as tmpfile: - fm.start_stream(tmpfile, None) - fm.request(f) - fm.response(f) - fm.stop_stream() - - assert read()[0].response - - with open(p, "ab") as tmpfile: - f = tutils.tflow() - fm.start_stream(tmpfile, None) - fm.request(f) - fm.shutdown() - - assert not read()[1].response - class TestRequest: |