improve mitmproxy.scripts semantics, clean up tests

1 files changed, 10 insertions, 2 deletions

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)