1 files changed, 45 insertions, 49 deletions
diff --git a/test/mitmproxy/addons/test_setheaders.py b/test/mitmproxy/addons/test_setheaders.py
index d7bdef61..34395ddf 100644
--- a/test/mitmproxy/addons/test_setheaders.py
+++ b/test/mitmproxy/addons/test_setheaders.py
@@ -1,21 +1,12 @@
 from mitmproxy.test import tflow
 from mitmproxy.test import tutils
-
-from .. import mastertest
+from mitmproxy.test import taddons
 
 from mitmproxy.addons import setheaders
 from mitmproxy import options
-from mitmproxy import proxy
-
 
-class TestSetHeaders(mastertest.MasterTest):
-    def mkmaster(self, **opts):
-        o = options.Options(**opts)
-        m = mastertest.RecordingMaster(o, proxy.DummyServer())
-        sh = setheaders.SetHeaders()
-        m.addons.add(sh)
-        return m, sh
 
+class TestSetHeaders:
     def test_configure(self):
         sh = setheaders.SetHeaders()
         o = options.Options(
@@ -27,41 +18,46 @@ class TestSetHeaders(mastertest.MasterTest):
         )
 
     def test_setheaders(self):
-        m, sh = self.mkmaster(
-            setheaders = [
-                ("~q", "one", "two"),
-                ("~s", "one", "three")
-            ]
-        )
-        f = tflow.tflow()
-        f.request.headers["one"] = "xxx"
-        m.request(f)
-        assert f.request.headers["one"] == "two"
-
-        f = tflow.tflow(resp=True)
-        f.response.headers["one"] = "xxx"
-        m.response(f)
-        assert f.response.headers["one"] == "three"
-
-        m, sh = self.mkmaster(
-            setheaders = [
-                ("~s", "one", "two"),
-                ("~s", "one", "three")
-            ]
-        )
-        f = tflow.tflow(resp=True)
-        f.request.headers["one"] = "xxx"
-        f.response.headers["one"] = "xxx"
-        m.response(f)
-        assert f.response.headers.get_all("one") == ["two", "three"]
-
-        m, sh = self.mkmaster(
-            setheaders = [
-                ("~q", "one", "two"),
-                ("~q", "one", "three")
-            ]
-        )
-        f = tflow.tflow()
-        f.request.headers["one"] = "xxx"
-        m.request(f)
-        assert f.request.headers.get_all("one") == ["two", "three"]
+        sh = setheaders.SetHeaders()
+        with taddons.context() as tctx:
+            tctx.configure(
+                sh,
+                setheaders = [
+                    ("~q", "one", "two"),
+                    ("~s", "one", "three")
+                ]
+            )
+            f = tflow.tflow()
+            f.request.headers["one"] = "xxx"
+            sh.request(f)
+            assert f.request.headers["one"] == "two"
+
+            f = tflow.tflow(resp=True)
+            f.response.headers["one"] = "xxx"
+            sh.response(f)
+            assert f.response.headers["one"] == "three"
+
+            tctx.configure(
+                sh,
+                setheaders = [
+                    ("~s", "one", "two"),
+                    ("~s", "one", "three")
+                ]
+            )
+            f = tflow.tflow(resp=True)
+            f.request.headers["one"] = "xxx"
+            f.response.headers["one"] = "xxx"
+            sh.response(f)
+            assert f.response.headers.get_all("one") == ["two", "three"]
+
+            tctx.configure(
+                sh,
+                setheaders = [
+                    ("~q", "one", "two"),
+                    ("~q", "one", "three")
+                ]
+            )
+            f = tflow.tflow()
+            f.request.headers["one"] = "xxx"
+            sh.request(f)
+            assert f.request.headers.get_all("one") == ["two", "three"]