1 files changed, 51 insertions, 0 deletions
diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py
new file mode 100644
index 00000000..270df207
--- /dev/null
+++ b/mitmproxy/test/taddons.py
@@ -0,0 +1,51 @@
+import mitmproxy.master
+import mitmproxy.options
+from mitmproxy import proxy
+from mitmproxy import events
+
+
+class context:
+    """
+        A context for testing addons, which sets up the mitmproxy.ctx module so
+        handlers can run as they would within mitmproxy. The context also
+        provides a number of helper methods for common testing scenarios.
+    """
+    def __init__(self, master = None, options = None):
+        self.options = options or mitmproxy.options.Options()
+        self.master = master or mitmproxy.master.Master(
+            options, proxy.DummyServer(options)
+        )
+        self.wrapped = None
+
+    def __enter__(self):
+        self.wrapped = self.master.handlecontext()
+        self.wrapped.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.wrapped.__exit__(exc_type, exc_value, traceback)
+        self.wrapped = None
+        return False
+
+    def cycle(self, addon, f):
+        """
+            Cycles the flow through the events for the flow. Stops if a reply
+            is taken (as in flow interception).
+        """
+        f.reply._state = "handled"
+        for evt, arg in events.event_sequence(f):
+            h = getattr(addon, evt, None)
+            if h:
+                h(arg)
+                if f.reply.state == "taken":
+                    return
+
+    def configure(self, addon, **kwargs):
+        """
+            A helper for testing configure methods. Modifies the registered
+            Options object with the given keyword arguments, then calls the
+            configure method on the addon with the updated value.
+        """
+        for k, v in kwargs.items():
+            setattr(self.options, k, v)
+        addon.configure(self.options, kwargs.keys())