Merge pull request #2051 from MatthewShao/fix-#1928-@concurrent-in-OOP

Fix #1928, @concurrent decorator in class

3 files changed, 36 insertions, 1 deletions

diff --git a/mitmproxy/script/concurrent.py b/mitmproxy/script/concurrent.py
index 366929a5..7573f2a5 100644
--- a/mitmproxy/script/concurrent.py
+++ b/mitmproxy/script/concurrent.py
@@ -29,4 +29,8 @@ def concurrent(fn):
             "script.concurrent (%s)" % fn.__name__,
             target=run
         ).start()
-    return _concurrent
+    # Support @concurrent for class-based addons
+    if "." in fn.__qualname__:
+        return staticmethod(_concurrent)
+    else:
+        return _concurrent
diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py
new file mode 100644
index 00000000..bd047c99
--- /dev/null
+++ b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py
@@ -0,0 +1,13 @@
+import time
+from mitmproxy.script import concurrent
+
+
+class ConcurrentClass:
+
+    @concurrent
+    def request(flow):
+        time.sleep(0.1)
+
+
+def start():
+    return ConcurrentClass()
diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py
index fb932d9a..e81c023d 100644
--- a/test/mitmproxy/script/test_concurrent.py
+++ b/test/mitmproxy/script/test_concurrent.py
@@ -44,3 +44,21 @@ class TestConcurrent(tservers.MasterTest):
             )
             sc.start()
             assert "decorator not supported" in tctx.master.event_log[0][1]
+
+    def test_concurrent_class(self):
+            with taddons.context() as tctx:
+                sc = script.Script(
+                    tutils.test_data.path(
+                        "mitmproxy/data/addonscripts/concurrent_decorator_class.py"
+                    )
+                )
+                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")