diff options
Diffstat (limited to 'test')
-rw-r--r-- | test/scripts/all.py | 4 | ||||
-rw-r--r-- | test/scripts/concurrent_decorator.py | 32 | ||||
-rw-r--r-- | test/scripts/concurrent_decorator_err.py | 5 | ||||
-rw-r--r-- | test/test_dump.py | 4 | ||||
-rw-r--r-- | test/test_flow.py | 6 | ||||
-rw-r--r-- | test/test_script.py | 53 | ||||
-rw-r--r-- | test/tutils.py | 19 |
7 files changed, 114 insertions, 9 deletions
diff --git a/test/scripts/all.py b/test/scripts/all.py index e6da7e51..7d30d757 100644 --- a/test/scripts/all.py +++ b/test/scripts/all.py @@ -3,6 +3,10 @@ def clientconnect(ctx, cc): ctx.log("XCLIENTCONNECT") log.append("clientconnect") +def serverconnect(ctx, cc): + ctx.log("XSERVERCONNECT") + log.append("serverconnect") + def request(ctx, r): ctx.log("XREQUEST") log.append("request") diff --git a/test/scripts/concurrent_decorator.py b/test/scripts/concurrent_decorator.py new file mode 100644 index 00000000..8e132006 --- /dev/null +++ b/test/scripts/concurrent_decorator.py @@ -0,0 +1,32 @@ +import time +from libmproxy.script import concurrent + + +@concurrent +def clientconnect(context, cc): + context.log("clientconnect") + + +@concurrent +def serverconnect(context, sc): + context.log("serverconnect") + + +@concurrent +def request(context, flow): + time.sleep(0.1) + + +@concurrent +def response(context, flow): + context.log("response") + + +@concurrent +def error(context, err): + context.log("error") + + +@concurrent +def clientdisconnect(context, dc): + context.log("clientdisconnect")
\ No newline at end of file diff --git a/test/scripts/concurrent_decorator_err.py b/test/scripts/concurrent_decorator_err.py new file mode 100644 index 00000000..78191315 --- /dev/null +++ b/test/scripts/concurrent_decorator_err.py @@ -0,0 +1,5 @@ +from libmproxy.script import concurrent + +@concurrent +def start(context, argv): + pass
\ No newline at end of file diff --git a/test/test_dump.py b/test/test_dump.py index 3b79c721..3d375f16 100644 --- a/test/test_dump.py +++ b/test/test_dump.py @@ -30,6 +30,9 @@ class TestDumpMaster: resp = tutils.tresp(req) resp.content = content m.handle_clientconnect(cc) + sc = proxy.ServerConnection(m.o, req.scheme, req.host, req.port, None) + sc.reply = mock.MagicMock() + m.handle_serverconnection(sc) m.handle_request(req) f = m.handle_response(resp) cd = flow.ClientDisconnect(cc) @@ -153,6 +156,7 @@ class TestDumpMaster: scripts=[[tutils.test_data.path("scripts/all.py")]], verbosity=0, eventlog=True ) assert "XCLIENTCONNECT" in ret + assert "XSERVERCONNECT" in ret assert "XREQUEST" in ret assert "XRESPONSE" in ret assert "XCLIENTDISCONNECT" in ret diff --git a/test/test_flow.py b/test/test_flow.py index 9844e0fd..c614960b 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -1,7 +1,7 @@ import Queue, time, os.path from cStringIO import StringIO import email.utils -from libmproxy import filt, flow, controller, utils, tnetstring +from libmproxy import filt, flow, controller, utils, tnetstring, proxy import tutils @@ -575,6 +575,10 @@ class TestFlowMaster: req = tutils.treq() fm.handle_clientconnect(req.client_conn) assert fm.scripts[0].ns["log"][-1] == "clientconnect" + sc = proxy.ServerConnection(None, req.scheme, req.host, req.port, None) + sc.reply = controller.DummyReply() + fm.handle_serverconnection(sc) + assert fm.scripts[0].ns["log"][-1] == "serverconnect" f = fm.handle_request(req) assert fm.scripts[0].ns["log"][-1] == "request" resp = tutils.tresp(req) diff --git a/test/test_script.py b/test/test_script.py index 9033c4fc..ad2296ef 100644 --- a/test/test_script.py +++ b/test/test_script.py @@ -2,6 +2,19 @@ from libmproxy import script, flow import tutils import shlex import os +import time + + +class TCounter: + count = 0 + + def __call__(self, *args, **kwargs): + self.count += 1 + + +class TScriptContext(TCounter): + def log(self, msg): + self.__call__() class TestScript: def test_simple(self): @@ -64,3 +77,43 @@ class TestScript: s.load ) + def test_concurrent(self): + s = flow.State() + fm = flow.FlowMaster(None, s) + fm.load_script([tutils.test_data.path("scripts/concurrent_decorator.py")]) + + reply = TCounter() + r1, r2 = tutils.treq(), tutils.treq() + r1.reply, r2.reply = reply, reply + t_start = time.time() + fm.handle_request(r1) + r1.reply() + fm.handle_request(r2) + r2.reply() + assert reply.count < 2 + assert (time.time() - t_start) < 0.09 + time.sleep(0.2) + assert reply.count == 2 + + def test_concurrent2(self): + ctx = TScriptContext() + s = script.Script([tutils.test_data.path("scripts/concurrent_decorator.py")], ctx) + s.load() + f = tutils.tflow_full() + f.error = tutils.terr(f.request) + f.reply = f.request.reply + + s.run("clientconnect", f) + s.run("serverconnect", f) + s.run("response", f) + s.run("error", f) + s.run("clientdisconnect", f) + time.sleep(0.1) + assert ctx.count == 5 + + def test_concurrent_err(self): + s = script.Script([tutils.test_data.path("scripts/concurrent_decorator_err.py")], TScriptContext()) + tutils.raises( + "decorator not supported for this method", + s.load + )
\ No newline at end of file diff --git a/test/tutils.py b/test/tutils.py index e42256ed..4cd7b7f8 100644 --- a/test/tutils.py +++ b/test/tutils.py @@ -33,6 +33,13 @@ def tresp(req=None): resp.reply = controller.DummyReply() return resp +def terr(req=None): + if not req: + req = treq() + err = flow.Error(req, "error") + err.reply = controller.DummyReply() + return err + def tflow(): r = treq() @@ -40,21 +47,17 @@ def tflow(): def tflow_full(): - r = treq() - f = flow.Flow(r) - f.response = tresp(r) + f = tflow() + f.response = tresp(f.request) return f def tflow_err(): - r = treq() - f = flow.Flow(r) - f.error = flow.Error(r, "error") - f.error.reply = controller.DummyReply() + f = tflow() + f.error = terr(f.request) return f - @contextmanager def tmpdir(*args, **kwargs): orig_workdir = os.getcwd() |