Merge pull request #1043 from mitmproxy/better-scripts

Better scripts

12 files changed, 186 insertions, 161 deletions

diff --git a/mitmproxy/console/__init__.py b/mitmproxy/console/__init__.py
index 5669be48..2f52d0b8 100644
--- a/mitmproxy/console/__init__.py
+++ b/mitmproxy/console/__init__.py
@@ -318,6 +318,7 @@ class ConsoleMaster(flow.FlowMaster):
 
         try:
             s = script.Script(command, script.ScriptContext(self))
+            s.load()
         except script.ScriptException as v:
             signals.status_message.send(
                 message = "Error loading script."
diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index 8fa84ed8..7de9823d 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -685,6 +685,7 @@ class FlowMaster(controller.Master):
         """
         try:
             s = script.Script(command, script.ScriptContext(self))
+            s.load()
         except script.ScriptException as v:
             return v.args[0]
         if use_reloader:
diff --git a/mitmproxy/script/reloader.py b/mitmproxy/script/reloader.py
index b4acf51b..99ce7f60 100644
--- a/mitmproxy/script/reloader.py
+++ b/mitmproxy/script/reloader.py
@@ -1,7 +1,7 @@
 import os
 import sys
 from watchdog.events import RegexMatchingEventHandler
-if sys.platform == 'darwin':
+if sys.platform == 'darwin':  # pragma: no cover
     from watchdog.observers.polling import PollingObserver as Observer
 else:
     from watchdog.observers import Observer
@@ -14,8 +14,8 @@ _observers = {}
 def watch(script, callback):
     if script in _observers:
         raise RuntimeError("Script already observed")
-    script_dir = os.path.dirname(os.path.abspath(script.args[0]))
-    script_name = os.path.basename(script.args[0])
+    script_dir = os.path.dirname(os.path.abspath(script.filename))
+    script_name = os.path.basename(script.filename)
     event_handler = _ScriptModificationHandler(callback, filename=script_name)
     observer = Observer()
     observer.schedule(event_handler, script_dir)
diff --git a/mitmproxy/script/script.py b/mitmproxy/script/script.py
index 55778851..edc17d43 100644
--- a/mitmproxy/script/script.py
+++ b/mitmproxy/script/script.py
@@ -22,7 +22,15 @@ class Script(object):
         self.args = self.parse_command(command)
         self.ctx = context
         self.ns = None
+
+    def __enter__(self):
         self.load()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_val:
+            return False  # reraise the exception
+        self.unload()
 
     @property
     def filename(self):
@@ -35,7 +43,7 @@ class Script(object):
         if os.name == "nt":  # Windows: escape all backslashes in the path.
             backslashes = shlex.split(command, posix=False)[0].count("\\")
             command = command.replace("\\", "\\\\", backslashes)
-        args = shlex.split(command)
+        args = shlex.split(command)  # pragma: nocover
         args[0] = os.path.expanduser(args[0])
         if not os.path.exists(args[0]):
             raise ScriptException(
@@ -58,7 +66,7 @@ class Script(object):
                 ScriptException on failure
         """
         if self.ns is not None:
-            self.unload()
+            raise ScriptException("Script is already loaded")
         script_dir = os.path.dirname(os.path.abspath(self.args[0]))
         self.ns = {'__file__': os.path.abspath(self.args[0])}
         sys.path.append(script_dir)
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)