Merge pull request #1132 from mhils/improve-script-loading

Improve script loading

3 files changed, 20 insertions, 18 deletions

diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py
index b560d9a1..c401a6b9 100644
--- a/test/mitmproxy/test_examples.py
+++ b/test/mitmproxy/test_examples.py
@@ -106,8 +106,8 @@ def test_modify_querystring():
 
 def test_modify_response_body():
     with tutils.raises(script.ScriptException):
-        with example("modify_response_body.py") as ex:
-            pass
+        with example("modify_response_body.py"):
+            assert True
 
     flow = tutils.tflow(resp=netutils.tresp(content="I <3 mitmproxy"))
     with example("modify_response_body.py mitmproxy rocks") as ex:
@@ -125,7 +125,7 @@ def test_redirect_requests():
 
 def test_har_extractor():
     with tutils.raises(script.ScriptException):
-        with example("har_extractor.py") as ex:
+        with example("har_extractor.py"):
             pass
 
     times = dict(
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 145e91cf..84a25689 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -7,7 +7,7 @@ import netlib.utils
 from netlib import odict
 from netlib.http import Headers
 from mitmproxy import filt, controller, tnetstring, flow
-from mitmproxy.exceptions import FlowReadException
+from mitmproxy.exceptions import FlowReadException, ScriptException
 from mitmproxy.models import Error
 from mitmproxy.models import Flow
 from mitmproxy.models import HTTPFlow
@@ -747,12 +747,16 @@ class TestFlowMaster:
     def test_load_script(self):
         s = flow.State()
         fm = flow.FlowMaster(None, s)
-        assert not fm.load_script(tutils.test_data.path("scripts/a.py"))
-        assert not fm.load_script(tutils.test_data.path("scripts/a.py"))
-        assert not fm.unload_scripts()
-        assert fm.load_script("nonexistent")
-        assert "ValueError" in fm.load_script(
-            tutils.test_data.path("scripts/starterr.py"))
+
+        fm.load_script(tutils.test_data.path("scripts/a.py"))
+        fm.load_script(tutils.test_data.path("scripts/a.py"))
+        fm.unload_scripts()
+        with tutils.raises(ScriptException):
+            fm.load_script("nonexistent")
+        try:
+            fm.load_script(tutils.test_data.path("scripts/starterr.py"))
+        except ScriptException as e:
+            assert "ValueError" in str(e)
         assert len(fm.scripts) == 0
 
     def test_getset_ignore(self):
@@ -779,7 +783,7 @@ class TestFlowMaster:
     def test_script_reqerr(self):
         s = flow.State()
         fm = flow.FlowMaster(None, s)
-        assert not fm.load_script(tutils.test_data.path("scripts/reqerr.py"))
+        fm.load_script(tutils.test_data.path("scripts/reqerr.py"))
         f = tutils.tflow()
         fm.handle_clientconnect(f.client_conn)
         assert fm.handle_request(f)
@@ -787,7 +791,7 @@ class TestFlowMaster:
     def test_script(self):
         s = flow.State()
         fm = flow.FlowMaster(None, s)
-        assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
+        fm.load_script(tutils.test_data.path("scripts/all.py"))
         f = tutils.tflow(resp=True)
 
         fm.handle_clientconnect(f.client_conn)
@@ -799,7 +803,7 @@ class TestFlowMaster:
         fm.handle_response(f)
         assert fm.scripts[0].ns["log"][-1] == "response"
         # load second script
-        assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
+        fm.load_script(tutils.test_data.path("scripts/all.py"))
         assert len(fm.scripts) == 2
         fm.handle_clientdisconnect(f.server_conn)
         assert fm.scripts[0].ns["log"][-1] == "clientdisconnect"
@@ -808,7 +812,7 @@ class TestFlowMaster:
         # unload first script
         fm.unload_scripts()
         assert len(fm.scripts) == 0
-        assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
+        fm.load_script(tutils.test_data.path("scripts/all.py"))
 
         f.error = tutils.terr()
         fm.handle_error(f)
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index 8843ee62..454736d4 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -285,8 +285,7 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin, AppMixin):
         self.master.set_stream_large_bodies(None)
 
     def test_stream_modify(self):
-        self.master.load_script(
-            tutils.test_data.path("scripts/stream_modify.py"))
+        self.master.load_script(tutils.test_data.path("scripts/stream_modify.py"))
         d = self.pathod('200:b"foo"')
         assert d.content == "bar"
         self.master.unload_scripts()
@@ -511,8 +510,7 @@ class TestTransparent(tservers.TransparentProxyTest, CommonMixin, TcpMixin):
     ssl = False
 
     def test_tcp_stream_modify(self):
-        self.master.load_script(
-            tutils.test_data.path("scripts/tcp_stream_modify.py"))
+        self.master.load_script(tutils.test_data.path("scripts/tcp_stream_modify.py"))
 
         self._tcpproxy_on()
         d = self.pathod('200:b"foo"')