aboutsummaryrefslogtreecommitdiffstats
path: root/test/mitmproxy/addons/test_readstdin.py
blob: 76b01f4fcc4a3c001df19a458c045f69baf26855 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
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