2 files changed, 64 insertions, 52 deletions
diff --git a/test/mitmproxy/builtins/test_setheaders.py b/test/mitmproxy/builtins/test_setheaders.py
new file mode 100644
index 00000000..1a8d048c
--- /dev/null
+++ b/test/mitmproxy/builtins/test_setheaders.py
@@ -0,0 +1,64 @@
+from .. import tutils, mastertest
+
+from mitmproxy.builtins import setheaders
+from mitmproxy.flow import state
+from mitmproxy.flow import options
+
+
+class TestSetHeaders(mastertest.MasterTest):
+    def mkmaster(self, **opts):
+        s = state.State()
+        m = mastertest.RecordingMaster(options.Options(**opts), None, s)
+        sh = setheaders.SetHeaders()
+        m.addons.add(sh)
+        return m, sh
+
+    def test_configure(self):
+        sh = setheaders.SetHeaders()
+        tutils.raises(
+            "invalid setheader filter pattern",
+            sh.configure,
+            options.Options(
+                setheaders = [("~b", "one", "two")]
+            )
+        )
+
+    def test_setheaders(self):
+        m, sh = self.mkmaster(
+            setheaders = [
+                ("~q", "one", "two"),
+                ("~s", "one", "three")
+            ]
+        )
+        f = tutils.tflow()
+        f.request.headers["one"] = "xxx"
+        self.invoke(m, "request", f)
+        assert f.request.headers["one"] == "two"
+
+        f = tutils.tflow(resp=True)
+        f.response.headers["one"] = "xxx"
+        self.invoke(m, "response", f)
+        assert f.response.headers["one"] == "three"
+
+        m, sh = self.mkmaster(
+            setheaders = [
+                ("~s", "one", "two"),
+                ("~s", "one", "three")
+            ]
+        )
+        f = tutils.tflow(resp=True)
+        f.request.headers["one"] = "xxx"
+        f.response.headers["one"] = "xxx"
+        self.invoke(m, "response", f)
+        assert f.response.headers.get_all("one") == ["two", "three"]
+
+        m, sh = self.mkmaster(
+            setheaders = [
+                ("~q", "one", "two"),
+                ("~q", "one", "three")
+            ]
+        )
+        f = tutils.tflow()
+        f.request.headers["one"] = "xxx"
+        self.invoke(m, "request", f)
+        assert f.request.headers.get_all("one") == ["two", "three"]
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 10163401..03935c16 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -961,55 +961,3 @@ class TestClientConnection:
         assert c3.get_state() == c.get_state()
 
         assert str(c)
-
-
-def test_setheaders():
-    h = flow.SetHeaders()
-    h.add("~q", "foo", "bar")
-    assert h.lst
-
-    h.set(
-        [
-            (".*", "one", "two"),
-            (".*", "three", "four"),
-        ]
-    )
-    assert h.count() == 2
-
-    h.clear()
-    assert not h.lst
-
-    h.add("~q", "foo", "bar")
-    h.add("~s", "foo", "bar")
-
-    v = h.get_specs()
-    assert v == [('~q', 'foo', 'bar'), ('~s', 'foo', 'bar')]
-    assert h.count() == 2
-    h.clear()
-    assert h.count() == 0
-
-    f = tutils.tflow()
-    f.request.content = b"foo"
-    h.add("~s", "foo", "bar")
-    h.run(f)
-    assert f.request.content == b"foo"
-
-    h.clear()
-    h.add("~s", "one", "two")
-    h.add("~s", "one", "three")
-    f = tutils.tflow(resp=True)
-    f.request.headers["one"] = "xxx"
-    f.response.headers["one"] = "xxx"
-    h.run(f)
-    assert f.request.headers["one"] == "xxx"
-    assert f.response.headers.get_all("one") == ["two", "three"]
-
-    h.clear()
-    h.add("~q", "one", "two")
-    h.add("~q", "one", "three")
-    f = tutils.tflow()
-    f.request.headers["one"] = "xxx"
-    h.run(f)
-    assert f.request.headers.get_all("one") == ["two", "three"]
-
-    assert not h.add("~", "foo", "bar")