test.taddon: Add cycle() method, use it to test addons.intercept

This commit is contained in:
Aldo Cortesi 2016-10-30 10:52:57 +13:00
parent b4904d33ba
commit d1f14961ed
4 changed files with 31 additions and 7 deletions

View File

@ -19,11 +19,10 @@ class Intercept:
def process_flow(self, f): def process_flow(self, f):
if self.filt: if self.filt:
should_intercept = all( should_intercept = all([
self.filt(f), self.filt(f),
not f.request.is_replay, not f.request.is_replay,
f.reply.state == "handled" ])
)
if should_intercept: if should_intercept:
f.intercept(ctx.master) f.intercept(ctx.master)

View File

@ -174,10 +174,10 @@ class Reply:
class DummyReply(Reply): class DummyReply(Reply):
""" """
A reply object that is not connected to anything. In contrast to regular Reply objects, A reply object that is not connected to anything. In contrast to regular
DummyReply objects are reset to "unhandled" at the end of an handler so that they can be used Reply objects, DummyReply objects are reset to "unhandled" at the end of an
multiple times. Useful when we need an object to seem like it has a channel, handler so that they can be used multiple times. Useful when we need an
and during testing. object to seem like it has a channel, and during testing.
""" """
def __init__(self): def __init__(self):
super().__init__(None) super().__init__(None)

View File

@ -1,6 +1,7 @@
import mitmproxy.master import mitmproxy.master
import mitmproxy.options import mitmproxy.options
from mitmproxy import proxy from mitmproxy import proxy
from mitmproxy import events
class context: class context:
@ -26,6 +27,19 @@ class context:
self.wrapped = None self.wrapped = None
return False return False
def cycle(self, addon, f):
"""
Cycles the flow through the events for the flow. Stops if a reply
is taken (as in flow interception).
"""
f.reply._state = "handled"
for evt, arg in events.event_sequence(f):
h = getattr(addon, evt, None)
if h:
h(arg)
if f.reply.state == "taken":
return
def configure(self, addon, **kwargs): def configure(self, addon, **kwargs):
""" """
A helper for testing configure methods. Modifies the registered A helper for testing configure methods. Modifies the registered

View File

@ -3,6 +3,7 @@ from mitmproxy import options
from mitmproxy import exceptions from mitmproxy import exceptions
from mitmproxy.test import taddons from mitmproxy.test import taddons
from mitmproxy.test import tutils from mitmproxy.test import tutils
from mitmproxy.test import tflow
class Options(options.Options): class Options(options.Options):
@ -23,3 +24,13 @@ def test_simple():
r, r,
intercept="~~" intercept="~~"
) )
tctx.configure(r, intercept="~s")
f = tflow.tflow(resp=True)
tctx.cycle(r, f)
assert f.intercepted
f = tflow.tflow(resp=False)
tctx.cycle(r, f)
assert not f.intercepted