diff options
Diffstat (limited to 'test')
| -rw-r--r-- | test/mitmproxy/addons/test_readfile.py | 93 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_readstdin.py | 53 |
2 files changed, 63 insertions, 83 deletions
diff --git a/test/mitmproxy/addons/test_readfile.py b/test/mitmproxy/addons/test_readfile.py index b30c147b..689d9779 100644 --- a/test/mitmproxy/addons/test_readfile.py +++ b/test/mitmproxy/addons/test_readfile.py @@ -1,62 +1,95 @@ +import io +from unittest import mock + +import pytest + +import mitmproxy.io +from mitmproxy import exceptions from mitmproxy.addons import readfile from mitmproxy.test import taddons from mitmproxy.test import tflow -from mitmproxy import io -from mitmproxy import exceptions -from unittest import mock -import pytest +@pytest.fixture +def data(): + f = io.BytesIO() -def write_data(path, corrupt=False): - with open(path, "wb") as tf: - w = io.FlowWriter(tf) - for i in range(3): - f = tflow.tflow(resp=True) - w.add(f) - for i in range(3): - f = tflow.tflow(err=True) - w.add(f) - f = tflow.ttcpflow() - w.add(f) - f = tflow.ttcpflow(err=True) - w.add(f) - if corrupt: - tf.write(b"flibble") + w = mitmproxy.io.FlowWriter(f) + flows = [ + tflow.tflow(resp=True), + tflow.tflow(err=True), + tflow.ttcpflow(), + tflow.ttcpflow(err=True) + ] + for flow in flows: + w.add(flow) + f.seek(0) + return f + + +@pytest.fixture +def corrupt_data(): + f = data() + f.seek(0, io.SEEK_END) + f.write(b"qibble") + f.seek(0) + return f -@mock.patch('mitmproxy.master.Master.load_flow') -def test_configure(mck, tmpdir): +@mock.patch('mitmproxy.master.Master.load_flow') +def test_configure(mck, tmpdir, data, corrupt_data): rf = readfile.ReadFile() with taddons.context() as tctx: - tf = str(tmpdir.join("tfile")) - write_data(tf) + tf = tmpdir.join("tfile") + + tf.write(data.getvalue()) tctx.configure(rf, rfile=str(tf)) assert not mck.called rf.running() assert mck.called - write_data(tf, corrupt=True) + tf.write(corrupt_data.getvalue()) tctx.configure(rf, rfile=str(tf)) with pytest.raises(exceptions.OptionsError): rf.running() @mock.patch('mitmproxy.master.Master.load_flow') -def test_corruption(mck, tmpdir): +@mock.patch('sys.stdin') +def test_configure_stdin(stdin, load_flow, data, corrupt_data): + rf = readfile.ReadFile() + with taddons.context() as tctx: + stdin.buffer = data + tctx.configure(rf, rfile="-") + assert not load_flow.called + rf.running() + assert load_flow.called + stdin.buffer = corrupt_data + tctx.configure(rf, rfile="-") + with pytest.raises(exceptions.OptionsError): + rf.running() + + +@mock.patch('mitmproxy.master.Master.load_flow') +def test_corrupt(mck, corrupt_data): rf = readfile.ReadFile() with taddons.context() as tctx: with pytest.raises(exceptions.FlowReadException): - rf.load_flows_file("nonexistent") + rf.load_flows(io.BytesIO(b"qibble")) assert not mck.called assert len(tctx.master.logs) == 1 - tfc = str(tmpdir.join("tfile")) - write_data(tfc, corrupt=True) - with pytest.raises(exceptions.FlowReadException): - rf.load_flows_file(tfc) + rf.load_flows(corrupt_data) assert mck.called assert len(tctx.master.logs) == 2 + + +def test_nonexisting_file(): + rf = readfile.ReadFile() + with taddons.context() as tctx: + with pytest.raises(exceptions.FlowReadException): + rf.load_flows_from_path("nonexistent") + assert len(tctx.master.logs) == 1 diff --git a/test/mitmproxy/addons/test_readstdin.py b/test/mitmproxy/addons/test_readstdin.py deleted file mode 100644 index 76b01f4f..00000000 --- a/test/mitmproxy/addons/test_readstdin.py +++ /dev/null @@ -1,53 +0,0 @@ - -import io -from mitmproxy.addons import readstdin -from mitmproxy.test import taddons -from mitmproxy.test import tflow -import mitmproxy.io -from unittest import mock - - -def gen_data(corrupt=False): - tf = io.BytesIO() - w = mitmproxy.io.FlowWriter(tf) - for i in range(3): - f = tflow.tflow(resp=True) - w.add(f) - for i in range(3): - f = tflow.tflow(err=True) - w.add(f) - f = tflow.ttcpflow() - w.add(f) - f = tflow.ttcpflow(err=True) - w.add(f) - if corrupt: - tf.write(b"flibble") - tf.seek(0) - return tf - - -class mStdin: - def __init__(self, d): - self.buffer = d - - def isatty(self): - return False - - -@mock.patch('mitmproxy.master.Master.load_flow') -def test_read(m, tmpdir): - rf = readstdin.ReadStdin() - with taddons.context() as tctx: - assert not m.called - rf.running(stdin=mStdin(gen_data())) - assert m.called - - rf.running(stdin=mStdin(None)) - assert tctx.master.logs - tctx.master.clear() - - m.reset_mock() - assert not m.called - rf.running(stdin=mStdin(gen_data(corrupt=True))) - assert m.called - assert tctx.master.logs |