Merge pull request #1714 from cortesi/justno2

Test and interface improvements

2 files changed, 24 insertions, 24 deletions

diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator.py b/test/mitmproxy/data/addonscripts/concurrent_decorator.py
index a56c2af1..162c00f4 100644
--- a/test/mitmproxy/data/addonscripts/concurrent_decorator.py
+++ b/test/mitmproxy/data/addonscripts/concurrent_decorator.py
@@ -1,6 +1,7 @@
 import time
 from mitmproxy.script import concurrent
 
+
 @concurrent
 def request(flow):
     time.sleep(0.1)
diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py
index 71b87d28..bb760f92 100644
--- a/test/mitmproxy/script/test_concurrent.py
+++ b/test/mitmproxy/script/test_concurrent.py
@@ -1,10 +1,9 @@
 from mitmproxy.test import tflow
 from mitmproxy.test import tutils
+from mitmproxy.test import taddons
+
 from mitmproxy import controller
 from mitmproxy.addons import script
-from mitmproxy import options
-from mitmproxy import proxy
-from mitmproxy import master
 
 import time
 
@@ -21,29 +20,29 @@ class Thing:
 class TestConcurrent(mastertest.MasterTest):
     @ttutils.skip_appveyor
     def test_concurrent(self):
-        m = master.Master(options.Options(), proxy.DummyServer())
-        sc = script.Script(
-            tutils.test_data.path(
-                "mitmproxy/data/addonscripts/concurrent_decorator.py"
+        with taddons.context() as tctx:
+            sc = script.Script(
+                tutils.test_data.path(
+                    "mitmproxy/data/addonscripts/concurrent_decorator.py"
+                )
             )
-        )
-        m.addons.add(sc)
-        f1, f2 = tflow.tflow(), tflow.tflow()
-        m.request(f1)
-        m.request(f2)
-        start = time.time()
-        while time.time() - start < 5:
-            if f1.reply.state == f2.reply.state == "committed":
-                return
-        raise ValueError("Script never acked")
+            sc.start()
+
+            f1, f2 = tflow.tflow(), tflow.tflow()
+            tctx.cycle(sc, f1)
+            tctx.cycle(sc, f2)
+            start = time.time()
+            while time.time() - start < 5:
+                if f1.reply.state == f2.reply.state == "committed":
+                    return
+            raise ValueError("Script never acked")
 
     def test_concurrent_err(self):
-        m = mastertest.RecordingMaster(options.Options(), proxy.DummyServer())
-        sc = script.Script(
-            tutils.test_data.path(
-                "mitmproxy/data/addonscripts/concurrent_decorator_err.py"
+        with taddons.context() as tctx:
+            sc = script.Script(
+                tutils.test_data.path(
+                    "mitmproxy/data/addonscripts/concurrent_decorator_err.py"
+                )
             )
-        )
-        with m.handlecontext():
             sc.start()
-        assert "decorator not supported" in m.event_log[0][1]
+            assert "decorator not supported" in tctx.master.event_log[0][1]