Use the contextmanager in all tests
This commit is contained in:
Shadab Zafar
parent
af9442a9ba
commit
7c5791e51a
|
|
@ -2,7 +2,7 @@ import glob
|
|||
import os
|
||||
from contextlib import contextmanager
|
||||
|
||||
from mitmproxy import utils, script
|
||||
from mitmproxy import utils, script, contentviews
|
||||
from mitmproxy.proxy import config
|
||||
from netlib import tutils as netutils
|
||||
from netlib.http import Headers
|
||||
|
|
@ -11,23 +11,24 @@ from . import tservers, tutils
|
|||
example_dir = utils.Data(__name__).path("../../examples")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def example(command):
|
||||
command = os.path.join(example_dir, command)
|
||||
# tmaster = tservers.TestMaster(config.ProxyConfig())
|
||||
# ctx = script.ScriptContext(tmaster)
|
||||
ctx = DummyContext()
|
||||
s = script.Script(command, ctx)
|
||||
yield s
|
||||
s.unload()
|
||||
|
||||
|
||||
class DummyContext(object):
|
||||
"""Emulate script.ScriptContext() functionality."""
|
||||
|
||||
def log(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def add_contentview(self, view_obj):
|
||||
pass
|
||||
|
||||
|
||||
@contextmanager
|
||||
def example(command):
|
||||
command = os.path.join(example_dir, command)
|
||||
ctx = DummyContext()
|
||||
s = script.Script(command, ctx)
|
||||
yield s
|
||||
s.unload()
|
||||
|
||||
|
||||
def test_load_scripts():
|
||||
scripts = glob.glob("%s/*.py" % example_dir)
|
||||
|
|
@ -56,55 +57,60 @@ def test_load_scripts():
|
|||
|
||||
def test_add_header():
|
||||
flow = tutils.tflow(resp=netutils.tresp())
|
||||
add_header.response({}, flow)
|
||||
assert flow.response.headers["newheader"] == "foo"
|
||||
with example("add_header.py") as ex:
|
||||
ex.run("response", flow)
|
||||
assert flow.response.headers["newheader"] == "foo"
|
||||
|
||||
|
||||
def test_custom_contentviews():
|
||||
with example("custom_contentviews.py"):
|
||||
pig = contentviews.get_by_shortcut("l")
|
||||
_, fmt = pig("<html>test!</html>")
|
||||
assert any('esttay!' in val[0][1] for val in fmt)
|
||||
assert not pig("gobbledygook")
|
||||
|
||||
|
||||
def test_iframe_injector():
|
||||
with tutils.raises(script.ScriptException):
|
||||
with example("iframe_injector.py") as ex:
|
||||
pass
|
||||
|
||||
flow = tutils.tflow(resp=netutils.tresp(content="<html>mitmproxy</html>"))
|
||||
with example("iframe_injector.py http://example.org/evil_iframe") as ex:
|
||||
ex.run("response", flow)
|
||||
content = flow.response.content
|
||||
assert 'iframe' in content and 'evil_iframe' in content
|
||||
|
||||
|
||||
def test_modify_form():
|
||||
form_header = Headers(content_type="application/x-www-form-urlencoded")
|
||||
flow = tutils.tflow(req=netutils.treq(headers=form_header))
|
||||
modify_form.request({}, flow)
|
||||
assert flow.request.urlencoded_form["mitmproxy"] == ["rocks"]
|
||||
with example("modify_form.py") as ex:
|
||||
ex.run("request", flow)
|
||||
assert flow.request.urlencoded_form["mitmproxy"] == ["rocks"]
|
||||
|
||||
|
||||
def test_modify_querystring():
|
||||
flow = tutils.tflow(req=netutils.treq(path="/search?q=term"))
|
||||
modify_querystring.request({}, flow)
|
||||
assert flow.request.query["mitmproxy"] == ["rocks"]
|
||||
with example("modify_querystring.py") as ex:
|
||||
ex.run("request", flow)
|
||||
assert flow.request.query["mitmproxy"] == ["rocks"]
|
||||
|
||||
|
||||
def test_modify_response_body():
|
||||
ctx = DummyContext()
|
||||
tutils.raises(ValueError, modify_response_body.start, ctx, [])
|
||||
|
||||
modify_response_body.start(ctx, ["modify-response-body.py", "mitmproxy", "rocks"])
|
||||
assert ctx.old == "mitmproxy" and ctx.new == "rocks"
|
||||
with tutils.raises(script.ScriptException):
|
||||
with example("modify_response_body.py") as ex:
|
||||
pass
|
||||
|
||||
flow = tutils.tflow(resp=netutils.tresp(content="I <3 mitmproxy"))
|
||||
modify_response_body.response(ctx, flow)
|
||||
assert flow.response.content == "I <3 rocks"
|
||||
|
||||
|
||||
def test_custom_contentviews():
|
||||
pig = custom_contentviews.ViewPigLatin()
|
||||
_, fmt = pig("<html>test!</html>")
|
||||
assert any('esttay!' in val[0][1] for val in fmt)
|
||||
assert not pig("gobbledygook")
|
||||
|
||||
|
||||
def test_iframe_injector():
|
||||
ctx = DummyContext()
|
||||
tutils.raises(ValueError, iframe_injector.start, ctx, [])
|
||||
|
||||
flow = tutils.tflow(resp=netutils.tresp(content="<html>Kungfu Panda 3</html>"))
|
||||
ctx.iframe_url = "http://example.org/evil_iframe"
|
||||
iframe_injector.response(ctx, flow)
|
||||
|
||||
content = flow.response.content
|
||||
assert 'iframe' in content and ctx.iframe_url in content
|
||||
with example("modify_response_body.py mitmproxy rocks") as ex:
|
||||
assert ex.ctx.old == "mitmproxy" and ex.ctx.new == "rocks"
|
||||
ex.run("response", flow)
|
||||
assert flow.response.content == "I <3 rocks"
|
||||
|
||||
|
||||
def test_redirect_requests():
|
||||
flow = tutils.tflow(req=netutils.treq(host="example.org"))
|
||||
redirect_requests.request({}, flow)
|
||||
assert flow.request.host == "mitmproxy.org"
|
||||
with example("redirect_requests.py") as ex:
|
||||
ex.run("request", flow)
|
||||
assert flow.request.host == "mitmproxy.org"
|
||||
|
|
|
|||
Loading…
Reference in New Issue