Merge pull request #1043 from mitmproxy/better-scripts

Better scripts

8 files changed, 171 insertions, 156 deletions

diff --git a/test/mitmproxy/script/__init__.py b/test/mitmproxy/script/__init__.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/test/mitmproxy/script/__init__.py
diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py
new file mode 100644
index 00000000..a9c9e153
--- /dev/null
+++ b/test/mitmproxy/script/test_concurrent.py
@@ -0,0 +1,32 @@
+from threading import Event
+
+from mitmproxy.script import Script
+from test.mitmproxy import tutils
+
+
+class Dummy:
+    def __init__(self, reply):
+        self.reply = reply
+
+
+@tutils.skip_appveyor
+def test_concurrent():
+    with Script(tutils.test_data.path("scripts/concurrent_decorator.py"), None) as s:
+        def reply():
+            reply.acked.set()
+        reply.acked = Event()
+
+        f1, f2 = Dummy(reply), Dummy(reply)
+        s.run("request", f1)
+        f1.reply()
+        s.run("request", f2)
+        f2.reply()
+        assert f1.reply.acked == reply.acked
+        assert not reply.acked.is_set()
+        assert reply.acked.wait(10)
+
+
+def test_concurrent_err():
+    s = Script(tutils.test_data.path("scripts/concurrent_decorator_err.py"), None)
+    with tutils.raises("Concurrent decorator not supported for 'start' method"):
+        s.load()
diff --git a/test/mitmproxy/script/test_reloader.py b/test/mitmproxy/script/test_reloader.py
new file mode 100644
index 00000000..0345f6ed
--- /dev/null
+++ b/test/mitmproxy/script/test_reloader.py
@@ -0,0 +1,34 @@
+import mock
+from mitmproxy.script.reloader import watch, unwatch
+from test.mitmproxy import tutils
+from threading import Event
+
+
+def test_simple():
+    with tutils.tmpdir():
+        with open("foo.py", "w"):
+            pass
+
+        script = mock.Mock()
+        script.filename = "foo.py"
+
+        e = Event()
+
+        def _onchange():
+            e.set()
+
+        watch(script, _onchange)
+        with tutils.raises("already observed"):
+            watch(script, _onchange)
+
+        # Some reloaders don't register a change directly after watching, because they first need to initialize.
+        # To test if watching works at all, we do repeated writes every 100ms.
+        for _ in range(100):
+            with open("foo.py", "a") as f:
+                f.write(".")
+            if e.wait(0.1):
+                break
+        else:
+            raise AssertionError("No change detected.")
+
+        unwatch(script)
diff --git a/test/mitmproxy/script/test_script.py b/test/mitmproxy/script/test_script.py
new file mode 100644
index 00000000..a9b55977
--- /dev/null
+++ b/test/mitmproxy/script/test_script.py
@@ -0,0 +1,83 @@
+from mitmproxy.script import Script
+from mitmproxy.exceptions import ScriptException
+from test.mitmproxy import tutils
+
+
+class TestParseCommand:
+    def test_empty_command(self):
+        with tutils.raises(ScriptException):
+            Script.parse_command("")
+
+        with tutils.raises(ScriptException):
+            Script.parse_command("  ")
+
+    def test_no_script_file(self):
+        with tutils.raises("not found"):
+            Script.parse_command("notfound")
+
+        with tutils.tmpdir() as dir:
+            with tutils.raises("not a file"):
+                Script.parse_command(dir)
+
+    def test_parse_args(self):
+        with tutils.chdir(tutils.test_data.dirname):
+            assert Script.parse_command("scripts/a.py") == ["scripts/a.py"]
+            assert Script.parse_command("scripts/a.py foo bar") == ["scripts/a.py", "foo", "bar"]
+            assert Script.parse_command("scripts/a.py 'foo bar'") == ["scripts/a.py", "foo bar"]
+
+    @tutils.skip_not_windows
+    def test_parse_windows(self):
+        with tutils.chdir(tutils.test_data.dirname):
+            assert Script.parse_command("scripts\\a.py") == ["scripts\\a.py"]
+            assert Script.parse_command("scripts\\a.py 'foo \\ bar'") == ["scripts\\a.py", 'foo \\ bar']
+
+
+def test_simple():
+    with tutils.chdir(tutils.test_data.path("scripts")):
+        s = Script("a.py --var 42", None)
+        assert s.filename == "a.py"
+        assert s.ns is None
+
+        s.load()
+        assert s.ns["var"] == 42
+
+        s.run("here")
+        assert s.ns["var"] == 43
+
+        s.unload()
+        assert s.ns is None
+
+        with tutils.raises(ScriptException):
+            s.run("here")
+
+        with Script("a.py --var 42", None) as s:
+            s.run("here")
+
+
+def test_script_exception():
+    with tutils.chdir(tutils.test_data.path("scripts")):
+        s = Script("syntaxerr.py", None)
+        with tutils.raises(ScriptException):
+            s.load()
+
+        s = Script("starterr.py", None)
+        with tutils.raises(ScriptException):
+            s.load()
+
+        s = Script("a.py", None)
+        s.load()
+        with tutils.raises(ScriptException):
+            s.load()
+
+        s = Script("a.py", None)
+        with tutils.raises(ScriptException):
+            s.run("here")
+
+        with tutils.raises(ScriptException):
+            with Script("reqerr.py", None) as s:
+                s.run("request", None)
+
+        s = Script("unloaderr.py", None)
+        s.load()
+        with tutils.raises(ScriptException):
+            s.unload()
diff --git a/test/mitmproxy/scripts/concurrent_decorator.py b/test/mitmproxy/scripts/concurrent_decorator.py
index 6651c811..cf3ab512 100644
--- a/test/mitmproxy/scripts/concurrent_decorator.py
+++ b/test/mitmproxy/scripts/concurrent_decorator.py
@@ -1,32 +1,6 @@
 import time
 from mitmproxy.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")
diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py
index 803776ac..b560d9a1 100644
--- a/test/mitmproxy/test_examples.py
+++ b/test/mitmproxy/test_examples.py
@@ -31,9 +31,8 @@ class DummyContext(object):
 def example(command):
     command = os.path.join(example_dir, command)
     ctx = DummyContext()
-    s = script.Script(command, ctx)
-    yield s
-    s.unload()
+    with script.Script(command, ctx) as s:
+        yield s
 
 
 def test_load_scripts():
@@ -52,8 +51,10 @@ def test_load_scripts():
             f += " ~a"
         if "modify_response_body" in f:
             f += " foo bar"  # two arguments required
+
+        s = script.Script(f, script.ScriptContext(tmaster))
         try:
-            s = script.Script(f, script.ScriptContext(tmaster))  # Loads the script file.
+            s.load()
         except Exception as v:
             if "ImportError" not in str(v):
                 raise
diff --git a/test/mitmproxy/test_script.py b/test/mitmproxy/test_script.py
index b827c623..f321d15c 100644
--- a/test/mitmproxy/test_script.py
+++ b/test/mitmproxy/test_script.py
@@ -1,27 +1,7 @@
-import os
-import time
-import mock
-from mitmproxy import script, flow
+from mitmproxy import flow
 from . import tutils
 
 
-def test_simple():
-    s = flow.State()
-    fm = flow.FlowMaster(None, s)
-    sp = tutils.test_data.path("scripts/a.py")
-    p = script.Script("%s --var 40" % sp, script.ScriptContext(fm))
-
-    assert "here" in p.ns
-    assert p.run("here") == 41
-    assert p.run("here") == 42
-
-    tutils.raises(script.ScriptException, p.run, "errargs")
-
-    # Check reload
-    p.load()
-    assert p.run("here") == 41
-
-
 def test_duplicate_flow():
     s = flow.State()
     fm = flow.FlowMaster(None, s)
@@ -31,102 +11,3 @@ def test_duplicate_flow():
     assert fm.state.flow_count() == 2
     assert not fm.state.view[0].request.is_replay
     assert fm.state.view[1].request.is_replay
-
-
-def test_err():
-    s = flow.State()
-    fm = flow.FlowMaster(None, s)
-    sc = script.ScriptContext(fm)
-
-    tutils.raises(
-        "not found",
-        script.Script, "nonexistent", sc
-    )
-
-    tutils.raises(
-        "not a file",
-        script.Script, tutils.test_data.path("scripts"), sc
-    )
-
-    tutils.raises(
-        script.ScriptException,
-        script.Script, tutils.test_data.path("scripts/syntaxerr.py"), sc
-    )
-
-    tutils.raises(
-        script.ScriptException,
-        script.Script, tutils.test_data.path("scripts/loaderr.py"), sc
-    )
-
-    scr = script.Script(tutils.test_data.path("scripts/unloaderr.py"), sc)
-    tutils.raises(script.ScriptException, scr.unload)
-
-
-@tutils.skip_appveyor
-def test_concurrent():
-    s = flow.State()
-    fm = flow.FlowMaster(None, s)
-    fm.load_script(tutils.test_data.path("scripts/concurrent_decorator.py"))
-
-    with mock.patch("mitmproxy.controller.DummyReply.__call__") as m:
-        f1, f2 = tutils.tflow(), tutils.tflow()
-        t_start = time.time()
-        fm.handle_request(f1)
-        f1.reply()
-        fm.handle_request(f2)
-        f2.reply()
-
-        # Two instantiations
-        assert m.call_count == 0  # No calls yet.
-        assert (time.time() - t_start) < 0.1
-
-
-def test_concurrent2():
-    s = flow.State()
-    fm = flow.FlowMaster(None, s)
-    s = script.Script(
-        tutils.test_data.path("scripts/concurrent_decorator.py"),
-        script.ScriptContext(fm))
-    s.load()
-    m = mock.Mock()
-
-    class Dummy:
-
-        def __init__(self):
-            self.response = self
-            self.error = self
-            self.reply = m
-
-    t_start = time.time()
-
-    for hook in ("clientconnect",
-                 "serverconnect",
-                 "response",
-                 "error",
-                 "clientconnect"):
-        d = Dummy()
-        s.run(hook, d)
-        d.reply()
-    while (time.time() - t_start) < 20 and m.call_count <= 5:
-        if m.call_count == 5:
-            return
-        time.sleep(0.001)
-    assert False
-
-
-def test_concurrent_err():
-    s = flow.State()
-    fm = flow.FlowMaster(None, s)
-    tutils.raises(
-        "Concurrent decorator not supported for 'start' method",
-        script.Script,
-        tutils.test_data.path("scripts/concurrent_decorator_err.py"),
-        fm)
-
-
-def test_command_parsing():
-    s = flow.State()
-    fm = flow.FlowMaster(None, s)
-    absfilepath = os.path.normcase(tutils.test_data.path("scripts/a.py"))
-    s = script.Script(absfilepath, script.ScriptContext(fm))
-    assert os.path.isfile(s.args[0])
diff --git a/test/mitmproxy/tutils.py b/test/mitmproxy/tutils.py
index 0d65df71..791db6d9 100644
--- a/test/mitmproxy/tutils.py
+++ b/test/mitmproxy/tutils.py
@@ -26,6 +26,13 @@ def skip_windows(fn):
         return fn
 
 
+def skip_not_windows(fn):
+    if os.name == "nt":
+        return fn
+    else:
+        return _skip_windows
+
+
 def _skip_appveyor(*args):
     raise SkipTest("Skipped on AppVeyor.")
 
@@ -120,14 +127,17 @@ def get_body_line(last_displayed_body, line_nb):
 
 
 @contextmanager
+def chdir(dir):
+    orig_dir = os.getcwd()
+    os.chdir(dir)
+    yield
+    os.chdir(orig_dir)
+
+@contextmanager
 def tmpdir(*args, **kwargs):
-    orig_workdir = os.getcwd()
     temp_workdir = tempfile.mkdtemp(*args, **kwargs)
-    os.chdir(temp_workdir)
-
-    yield temp_workdir
-
-    os.chdir(orig_workdir)
+    with chdir(temp_workdir):
+        yield temp_workdir
     shutil.rmtree(temp_workdir)