aboutsummaryrefslogtreecommitdiffstats
path: root/test/mitmproxy/addons/test_readfile.py
blob: b30c147b4118b4b703d7481f84b12b2ad69518e4 (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
54
55
56
57
58
59
60
61
62
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


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")


@mock.patch('mitmproxy.master.Master.load_flow')
def test_configure(mck, tmpdir):

    rf = readfile.ReadFile()
    with taddons.context() as tctx:
        tf = str(tmpdir.join("tfile"))
        write_data(tf)
        tctx.configure(rf, rfile=str(tf))
        assert not mck.called
        rf.running()
        assert mck.called

        write_data(tf, corrupt=True)
        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):

    rf = readfile.ReadFile()
    with taddons.context() as tctx:
        with pytest.raises(exceptions.FlowReadException):
            rf.load_flows_file("nonexistent")
        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)
        assert mck.called
        assert len(tctx.master.logs) == 2