From 6c8c4465d9971076e9a04f8042aebd5d71f67d59 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Fri, 4 Nov 2016 09:32:06 +1300
Subject: [PATCH] script.concurrent: tests to taddons

---
 .../data/addonscripts/concurrent_decorator.py |  1 +
 test/mitmproxy/script/test_concurrent.py      | 47 +++++++++----------
 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 a56c2af1c..162c00f4b 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 71b87d288..bb760f926 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]