integrate readstdin into readfile addon

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