From 7c5791e51a55aebd909c950e8b265f76fca45d4c Mon Sep 17 00:00:00 2001 From: Shadab Zafar Date: Mon, 7 Mar 2016 22:53:12 +0530 Subject: [PATCH] Use the contextmanager in all tests --- examples/__init__.py | 0 test/mitmproxy/test_examples.py | 98 +++++++++++++++++---------------- 2 files changed, 52 insertions(+), 46 deletions(-) delete mode 100644 examples/__init__.py diff --git a/examples/__init__.py b/examples/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py index 3244ebf23..a6fed38be 100644 --- a/test/mitmproxy/test_examples.py +++ b/test/mitmproxy/test_examples.py @@ -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("test!") + 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="mitmproxy")) + 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("test!") - 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="Kungfu Panda 3")) - 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"