addons.setheaders: fix configure bug, tests to taddons

This commit is contained in:
Aldo Cortesi 2016-11-04 09:09:39 +13:00
parent b51a96081a
commit deb66d3cac
2 changed files with 51 additions and 53 deletions

View File

@ -14,13 +14,15 @@ class SetHeaders:
header: Header name.
value: Header value string
"""
for fpatt, header, value in options.setheaders:
flt = flowfilter.parse(fpatt)
if not flt:
raise exceptions.OptionsError(
"Invalid setheader filter pattern %s" % fpatt
)
self.lst.append((fpatt, header, value, flt))
if "setheaders" in updated:
self.lst = []
for fpatt, header, value in options.setheaders:
flt = flowfilter.parse(fpatt)
if not flt:
raise exceptions.OptionsError(
"Invalid setheader filter pattern %s" % fpatt
)
self.lst.append((fpatt, header, value, flt))
def run(self, f, hdrs):
for _, header, value, flt in self.lst:

View File

@ -1,21 +1,12 @@
from mitmproxy.test import tflow
from mitmproxy.test import tutils
from .. import mastertest
from mitmproxy.test import taddons
from mitmproxy.addons import setheaders
from mitmproxy import options
from mitmproxy import proxy
class TestSetHeaders(mastertest.MasterTest):
def mkmaster(self, **opts):
o = options.Options(**opts)
m = mastertest.RecordingMaster(o, proxy.DummyServer())
sh = setheaders.SetHeaders()
m.addons.add(sh)
return m, sh
class TestSetHeaders:
def test_configure(self):
sh = setheaders.SetHeaders()
o = options.Options(
@ -27,41 +18,46 @@ class TestSetHeaders(mastertest.MasterTest):
)
def test_setheaders(self):
m, sh = self.mkmaster(
setheaders = [
("~q", "one", "two"),
("~s", "one", "three")
]
)
f = tflow.tflow()
f.request.headers["one"] = "xxx"
m.request(f)
assert f.request.headers["one"] == "two"
sh = setheaders.SetHeaders()
with taddons.context() as tctx:
tctx.configure(
sh,
setheaders = [
("~q", "one", "two"),
("~s", "one", "three")
]
)
f = tflow.tflow()
f.request.headers["one"] = "xxx"
sh.request(f)
assert f.request.headers["one"] == "two"
f = tflow.tflow(resp=True)
f.response.headers["one"] = "xxx"
m.response(f)
assert f.response.headers["one"] == "three"
f = tflow.tflow(resp=True)
f.response.headers["one"] = "xxx"
sh.response(f)
assert f.response.headers["one"] == "three"
m, sh = self.mkmaster(
setheaders = [
("~s", "one", "two"),
("~s", "one", "three")
]
)
f = tflow.tflow(resp=True)
f.request.headers["one"] = "xxx"
f.response.headers["one"] = "xxx"
m.response(f)
assert f.response.headers.get_all("one") == ["two", "three"]
tctx.configure(
sh,
setheaders = [
("~s", "one", "two"),
("~s", "one", "three")
]
)
f = tflow.tflow(resp=True)
f.request.headers["one"] = "xxx"
f.response.headers["one"] = "xxx"
sh.response(f)
assert f.response.headers.get_all("one") == ["two", "three"]
m, sh = self.mkmaster(
setheaders = [
("~q", "one", "two"),
("~q", "one", "three")
]
)
f = tflow.tflow()
f.request.headers["one"] = "xxx"
m.request(f)
assert f.request.headers.get_all("one") == ["two", "three"]
tctx.configure(
sh,
setheaders = [
("~q", "one", "two"),
("~q", "one", "three")
]
)
f = tflow.tflow()
f.request.headers["one"] = "xxx"
sh.request(f)
assert f.request.headers.get_all("one") == ["two", "three"]