From deb66d3cac6106081d7c2260bdfed944a107a081 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Fri, 4 Nov 2016 09:09:39 +1300
Subject: [PATCH] addons.setheaders: fix configure bug, tests to taddons

---
 mitmproxy/addons/setheaders.py           | 16 +++--
 test/mitmproxy/addons/test_setheaders.py | 88 +++++++++++-------------
 2 files changed, 51 insertions(+), 53 deletions(-)

diff --git a/mitmproxy/addons/setheaders.py b/mitmproxy/addons/setheaders.py
index 5695e1e8a..601e75216 100644
--- a/mitmproxy/addons/setheaders.py
+++ b/mitmproxy/addons/setheaders.py
@@ -14,13 +14,15 @@ class SetHeaders:
             header: Header name.
             value: Header value string
         """
-        for fpatt, header, value in options.setheaders:
-            flt = flowfilter.parse(fpatt)
-            if not flt:
-                raise exceptions.OptionsError(
-                    "Invalid setheader filter pattern %s" % fpatt
-                )
-            self.lst.append((fpatt, header, value, flt))
+        if "setheaders" in updated:
+            self.lst = []
+            for fpatt, header, value in options.setheaders:
+                flt = flowfilter.parse(fpatt)
+                if not flt:
+                    raise exceptions.OptionsError(
+                        "Invalid setheader filter pattern %s" % fpatt
+                    )
+                self.lst.append((fpatt, header, value, flt))
 
     def run(self, f, hdrs):
         for _, header, value, flt in self.lst:
diff --git a/test/mitmproxy/addons/test_setheaders.py b/test/mitmproxy/addons/test_setheaders.py
index d7bdef61c..34395ddfc 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"
+        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"
-        m.response(f)
-        assert f.response.headers["one"] == "three"
+            f = tflow.tflow(resp=True)
+            f.response.headers["one"] = "xxx"
+            sh.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"]
+            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"]
 
-        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"]
+            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"]