Convert flow.state to an addon

A first step, and already many things are much nicer.

18 files changed, 81 insertions, 136 deletions

diff --git a/test/mitmproxy/builtins/test_anticache.py b/test/mitmproxy/builtins/test_anticache.py
index 81041e73..512c90f5 100644
--- a/test/mitmproxy/builtins/test_anticache.py
+++ b/test/mitmproxy/builtins/test_anticache.py
@@ -1,15 +1,13 @@
 from .. import tutils, mastertest
 from mitmproxy.builtins import anticache
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 from mitmproxy import options
 
 
 class TestAntiCache(mastertest.MasterTest):
     def test_simple(self):
-        s = state.State()
         o = options.Options(anticache = True)
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sa = anticache.AntiCache()
         m.addons.add(sa)
 
diff --git a/test/mitmproxy/builtins/test_anticomp.py b/test/mitmproxy/builtins/test_anticomp.py
index f866f33f..1014b9ba 100644
--- a/test/mitmproxy/builtins/test_anticomp.py
+++ b/test/mitmproxy/builtins/test_anticomp.py
@@ -1,15 +1,13 @@
 from .. import tutils, mastertest
 from mitmproxy.builtins import anticomp
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 from mitmproxy import options
 
 
 class TestAntiComp(mastertest.MasterTest):
     def test_simple(self):
-        s = state.State()
         o = options.Options(anticomp = True)
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sa = anticomp.AntiComp()
         m.addons.add(sa)
 
diff --git a/test/mitmproxy/builtins/test_dumper.py b/test/mitmproxy/builtins/test_dumper.py
index 20403eea..3ab75bee 100644
--- a/test/mitmproxy/builtins/test_dumper.py
+++ b/test/mitmproxy/builtins/test_dumper.py
@@ -3,7 +3,6 @@ import io
 from .. import tutils, mastertest
 
 from mitmproxy.builtins import dumper
-from mitmproxy.flow import state
 from mitmproxy import exceptions
 from mitmproxy import dump
 from mitmproxy import models
@@ -71,14 +70,13 @@ class TestContentView(mastertest.MasterTest):
     def test_contentview(self, view_auto):
         view_auto.side_effect = exceptions.ContentViewException("")
 
-        s = state.State()
         sio = io.StringIO()
         o = dump.Options(
             flow_detail=4,
             verbosity=3,
             tfile=sio,
         )
-        m = mastertest.RecordingMaster(o, None, s)
+        m = mastertest.RecordingMaster(o, None)
         d = dumper.Dumper()
         m.addons.add(d)
         m.response(tutils.tflow())
diff --git a/test/mitmproxy/builtins/test_filestreamer.py b/test/mitmproxy/builtins/test_filestreamer.py
index a43ea0b7..6fbeb40b 100644
--- a/test/mitmproxy/builtins/test_filestreamer.py
+++ b/test/mitmproxy/builtins/test_filestreamer.py
@@ -4,7 +4,6 @@ import os.path
 
 from mitmproxy.builtins import filestreamer
 from mitmproxy.flow import master, FlowReader
-from mitmproxy.flow import state
 from mitmproxy import options
 
 
@@ -17,11 +16,10 @@ class TestStream(mastertest.MasterTest):
                 r = FlowReader(open(p, "rb"))
                 return list(r.stream())
 
-            s = state.State()
             o = options.Options(
                 outfile = (p, "wb")
             )
-            m = master.FlowMaster(o, None, s)
+            m = master.FlowMaster(o, None)
             sa = filestreamer.FileStreamer()
 
             m.addons.add(sa)
diff --git a/test/mitmproxy/builtins/test_replace.py b/test/mitmproxy/builtins/test_replace.py
index e9869ab7..03943867 100644
--- a/test/mitmproxy/builtins/test_replace.py
+++ b/test/mitmproxy/builtins/test_replace.py
@@ -1,7 +1,6 @@
 from .. import tutils, mastertest
 from mitmproxy.builtins import replace
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 from mitmproxy import options
 
 
@@ -30,14 +29,13 @@ class TestReplace(mastertest.MasterTest):
         )
 
     def test_simple(self):
-        s = state.State()
         o = options.Options(
             replacements = [
                 ("~q", "foo", "bar"),
                 ("~s", "foo", "bar"),
             ]
         )
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sa = replace.Replace()
         m.addons.add(sa)
 
diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py
index 683f3064..ecf2f3f7 100644
--- a/test/mitmproxy/builtins/test_script.py
+++ b/test/mitmproxy/builtins/test_script.py
@@ -8,7 +8,6 @@ from mitmproxy import exceptions
 from mitmproxy import options
 from mitmproxy.builtins import script
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 
 from .. import tutils, mastertest
 
@@ -57,9 +56,8 @@ def test_load_script():
 
 class TestScript(mastertest.MasterTest):
     def test_simple(self):
-        s = state.State()
         o = options.Options()
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sc = script.Script(
             tutils.test_data.path(
                 "data/addonscripts/recorder.py"
@@ -79,9 +77,8 @@ class TestScript(mastertest.MasterTest):
         assert recf[1] == "request"
 
     def test_reload(self):
-        s = state.State()
         o = options.Options()
-        m = mastertest.RecordingMaster(o, None, s)
+        m = mastertest.RecordingMaster(o, None)
         with tutils.tmpdir():
             with open("foo.py", "w"):
                 pass
@@ -98,9 +95,8 @@ class TestScript(mastertest.MasterTest):
             raise AssertionError("Change event not detected.")
 
     def test_exception(self):
-        s = state.State()
         o = options.Options()
-        m = mastertest.RecordingMaster(o, None, s)
+        m = mastertest.RecordingMaster(o, None)
         sc = script.Script(
             tutils.test_data.path("data/addonscripts/error.py")
         )
@@ -113,25 +109,9 @@ class TestScript(mastertest.MasterTest):
         assert re.search('addonscripts/error.py", line \d+, in mkerr', m.event_log[0][1])
         assert m.event_log[0][1].endswith("ValueError: Error!\n")
 
-    def test_duplicate_flow(self):
-        s = state.State()
-        o = options.Options()
-        fm = master.FlowMaster(o, None, s)
-        fm.addons.add(
-            script.Script(
-                tutils.test_data.path("data/addonscripts/duplicate_flow.py")
-            )
-        )
-        f = tutils.tflow()
-        fm.request(f)
-        assert fm.state.flow_count() == 2
-        assert not fm.state.view[0].request.is_replay
-        assert fm.state.view[1].request.is_replay
-
     def test_addon(self):
-        s = state.State()
         o = options.Options()
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sc = script.Script(
             tutils.test_data.path(
                 "data/addonscripts/addon.py"
@@ -163,9 +143,8 @@ class TestCutTraceback:
 
 class TestScriptLoader(mastertest.MasterTest):
     def test_run_once(self):
-        s = state.State()
         o = options.Options(scripts=[])
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sl = script.ScriptLoader()
         m.addons.add(sl)
 
@@ -188,9 +167,8 @@ class TestScriptLoader(mastertest.MasterTest):
             )
 
     def test_simple(self):
-        s = state.State()
         o = options.Options(scripts=[])
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sc = script.ScriptLoader()
         m.addons.add(sc)
         assert len(m.addons) == 1
@@ -204,16 +182,14 @@ class TestScriptLoader(mastertest.MasterTest):
         assert len(m.addons) == 1
 
     def test_dupes(self):
-        s = state.State()
         o = options.Options(scripts=["one", "one"])
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sc = script.ScriptLoader()
         tutils.raises(exceptions.OptionsError, m.addons.add, o, sc)
 
     def test_order(self):
         rec = tutils.test_data.path("data/addonscripts/recorder.py")
 
-        s = state.State()
         o = options.Options(
             scripts = [
                 "%s %s" % (rec, "a"),
@@ -221,7 +197,7 @@ class TestScriptLoader(mastertest.MasterTest):
                 "%s %s" % (rec, "c"),
             ]
         )
-        m = mastertest.RecordingMaster(o, None, s)
+        m = mastertest.RecordingMaster(o, None)
         sc = script.ScriptLoader()
         m.addons.add(sc)
 
diff --git a/test/mitmproxy/builtins/test_serverplayback.py b/test/mitmproxy/builtins/test_serverplayback.py
index 55ce9350..b97c01dd 100644
--- a/test/mitmproxy/builtins/test_serverplayback.py
+++ b/test/mitmproxy/builtins/test_serverplayback.py
@@ -4,7 +4,6 @@ import netlib.tutils
 from mitmproxy.builtins import serverplayback
 from mitmproxy import options
 from mitmproxy import exceptions
-from mitmproxy import flow
 
 
 class TestServerPlayback:
@@ -238,10 +237,9 @@ class TestServerPlayback:
         assert not s._hash(r) == s._hash(r2)
 
     def test_server_playback_full(self):
-        state = flow.State()
         s = serverplayback.ServerPlayback()
         o = options.Options(refresh_server_playback = True, keepserving=False)
-        m = mastertest.RecordingMaster(o, None, state)
+        m = mastertest.RecordingMaster(o, None)
         m.addons.add(s)
 
         f = tutils.tflow()
@@ -268,10 +266,9 @@ class TestServerPlayback:
         assert s.stop
 
     def test_server_playback_kill(self):
-        state = flow.State()
         s = serverplayback.ServerPlayback()
         o = options.Options(refresh_server_playback = True, replay_kill_extra=True)
-        m = mastertest.RecordingMaster(o, None, state)
+        m = mastertest.RecordingMaster(o, None)
         m.addons.add(s)
 
         f = tutils.tflow()
diff --git a/test/mitmproxy/builtins/test_setheaders.py b/test/mitmproxy/builtins/test_setheaders.py
index d148234d..234341f4 100644
--- a/test/mitmproxy/builtins/test_setheaders.py
+++ b/test/mitmproxy/builtins/test_setheaders.py
@@ -1,15 +1,13 @@
 from .. import tutils, mastertest
 
 from mitmproxy.builtins import setheaders
-from mitmproxy.flow import state
 from mitmproxy import options
 
 
 class TestSetHeaders(mastertest.MasterTest):
     def mkmaster(self, **opts):
-        s = state.State()
         o = options.Options(**opts)
-        m = mastertest.RecordingMaster(o, None, s)
+        m = mastertest.RecordingMaster(o, None)
         sh = setheaders.SetHeaders()
         m.addons.add(sh)
         return m, sh
diff --git a/test/mitmproxy/builtins/test_stickyauth.py b/test/mitmproxy/builtins/test_stickyauth.py
index 8b939c24..25141554 100644
--- a/test/mitmproxy/builtins/test_stickyauth.py
+++ b/test/mitmproxy/builtins/test_stickyauth.py
@@ -1,15 +1,13 @@
 from .. import tutils, mastertest
 from mitmproxy.builtins import stickyauth
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 from mitmproxy import options
 
 
 class TestStickyAuth(mastertest.MasterTest):
     def test_simple(self):
-        s = state.State()
         o = options.Options(stickyauth = ".*")
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sa = stickyauth.StickyAuth()
         m.addons.add(sa)
 
diff --git a/test/mitmproxy/builtins/test_stickycookie.py b/test/mitmproxy/builtins/test_stickycookie.py
index 4d3b7cb3..28b05a11 100644
--- a/test/mitmproxy/builtins/test_stickycookie.py
+++ b/test/mitmproxy/builtins/test_stickycookie.py
@@ -1,7 +1,6 @@
 from .. import tutils, mastertest
 from mitmproxy.builtins import stickycookie
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 from mitmproxy import options
 from netlib import tutils as ntutils
 
@@ -13,12 +12,11 @@ def test_domain_match():
 
 class TestStickyCookie(mastertest.MasterTest):
     def mk(self):
-        s = state.State()
         o = options.Options(stickycookie = ".*")
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sc = stickycookie.StickyCookie()
         m.addons.add(sc)
-        return s, m, sc
+        return m, sc
 
     def test_config(self):
         sc = stickycookie.StickyCookie()
@@ -29,7 +27,7 @@ class TestStickyCookie(mastertest.MasterTest):
         )
 
     def test_simple(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
         m.addons.add(sc)
 
         f = tutils.tflow(resp=True)
@@ -47,44 +45,44 @@ class TestStickyCookie(mastertest.MasterTest):
         m.request(f)
         assert f.request.headers["cookie"] == "foo=bar"
 
-    def _response(self, s, m, sc, cookie, host):
+    def _response(self, m, sc, cookie, host):
         f = tutils.tflow(req=ntutils.treq(host=host, port=80), resp=True)
         f.response.headers["Set-Cookie"] = cookie
         m.response(f)
         return f
 
     def test_response(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
 
         c = "SSID=mooo; domain=.google.com, FOO=bar; Domain=.google.com; Path=/; " \
             "Expires=Wed, 13-Jan-2021 22:23:01 GMT; Secure; "
 
-        self._response(s, m, sc, c, "host")
+        self._response(m, sc, c, "host")
         assert not sc.jar.keys()
 
-        self._response(s, m, sc, c, "www.google.com")
+        self._response(m, sc, c, "www.google.com")
         assert sc.jar.keys()
 
         sc.jar.clear()
         self._response(
-            s, m, sc, "SSID=mooo", "www.google.com"
+            m, sc, "SSID=mooo", "www.google.com"
         )
         assert list(sc.jar.keys())[0] == ('www.google.com', 80, '/')
 
     def test_response_multiple(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
 
         # Test setting of multiple cookies
         c1 = "somecookie=test; Path=/"
         c2 = "othercookie=helloworld; Path=/"
-        f = self._response(s, m, sc, c1, "www.google.com")
+        f = self._response(m, sc, c1, "www.google.com")
         f.response.headers["Set-Cookie"] = c2
         m.response(f)
         googlekey = list(sc.jar.keys())[0]
         assert len(sc.jar[googlekey].keys()) == 2
 
     def test_response_weird(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
 
         # Test setting of weird cookie keys
         f = tutils.tflow(req=ntutils.treq(host="www.google.com", port=80), resp=True)
@@ -100,12 +98,12 @@ class TestStickyCookie(mastertest.MasterTest):
         assert len(sc.jar[googlekey].keys()) == len(cs)
 
     def test_response_overwrite(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
 
         # Test overwriting of a cookie value
         c1 = "somecookie=helloworld; Path=/"
         c2 = "somecookie=newvalue; Path=/"
-        f = self._response(s, m, sc, c1, "www.google.com")
+        f = self._response(m, sc, c1, "www.google.com")
         f.response.headers["Set-Cookie"] = c2
         m.response(f)
         googlekey = list(sc.jar.keys())[0]
@@ -113,19 +111,19 @@ class TestStickyCookie(mastertest.MasterTest):
         assert list(sc.jar[googlekey]["somecookie"].items())[0][1] == "newvalue"
 
     def test_response_delete(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
 
         # Test that a cookie is be deleted
         # by setting the expire time in the past
-        f = self._response(s, m, sc, "duffer=zafar; Path=/", "www.google.com")
+        f = self._response(m, sc, "duffer=zafar; Path=/", "www.google.com")
         f.response.headers["Set-Cookie"] = "duffer=; Expires=Thu, 01-Jan-1970 00:00:00 GMT"
         m.response(f)
         assert not sc.jar.keys()
 
     def test_request(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
 
-        f = self._response(s, m, sc, "SSID=mooo", "www.google.com")
+        f = self._response(m, sc, "SSID=mooo", "www.google.com")
         assert "cookie" not in f.request.headers
         m.request(f)
         assert "cookie" in f.request.headers
diff --git a/test/mitmproxy/builtins/test_streambodies.py b/test/mitmproxy/builtins/test_streambodies.py
index 0e8a82f2..0a306681 100644
--- a/test/mitmproxy/builtins/test_streambodies.py
+++ b/test/mitmproxy/builtins/test_streambodies.py
@@ -1,5 +1,4 @@
 from .. import tutils, mastertest
-from mitmproxy.flow import state
 from mitmproxy.flow import master
 from mitmproxy import options
 
@@ -8,9 +7,8 @@ from mitmproxy.builtins import streambodies
 
 class TestStreamBodies(mastertest.MasterTest):
     def test_simple(self):
-        s = state.DummyState()
         o = options.Options(stream_large_bodies = 10)
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sa = streambodies.StreamBodies()
         m.addons.add(sa)
 
diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py
index 5605b1a6..ae11e577 100644
--- a/test/mitmproxy/mastertest.py
+++ b/test/mitmproxy/mastertest.py
@@ -50,8 +50,7 @@ class RecordingMaster(master.FlowMaster):
 
 @contextlib.contextmanager
 def mockctx():
-    state = flow.State()
     o = options.Options(refresh_server_playback = True, keepserving=False)
-    m = RecordingMaster(o, proxy.DummyServer(o), state)
+    m = RecordingMaster(o, proxy.DummyServer(o))
     with m.handlecontext():
         yield
diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py
index 5f943aeb..eb0e83f7 100644
--- a/test/mitmproxy/script/test_concurrent.py
+++ b/test/mitmproxy/script/test_concurrent.py
@@ -3,7 +3,6 @@ from mitmproxy import controller
 from mitmproxy.builtins import script
 from mitmproxy import options
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 import time
 
 
@@ -16,8 +15,7 @@ class Thing:
 class TestConcurrent(mastertest.MasterTest):
     @tutils.skip_appveyor
     def test_concurrent(self):
-        s = state.State()
-        m = master.FlowMaster(options.Options(), None, s)
+        m = master.FlowMaster(options.Options(), None)
         sc = script.Script(
             tutils.test_data.path(
                 "data/addonscripts/concurrent_decorator.py"
@@ -34,8 +32,7 @@ class TestConcurrent(mastertest.MasterTest):
         raise ValueError("Script never acked")
 
     def test_concurrent_err(self):
-        s = state.State()
-        m = mastertest.RecordingMaster(options.Options(), None, s)
+        m = mastertest.RecordingMaster(options.Options(), None)
         sc = script.Script(
             tutils.test_data.path(
                 "data/addonscripts/concurrent_decorator_err.py"
diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py
index 05b7544f..f824584b 100644
--- a/test/mitmproxy/test_examples.py
+++ b/test/mitmproxy/test_examples.py
@@ -6,7 +6,6 @@ from mitmproxy import options
 from mitmproxy import contentviews
 from mitmproxy.builtins import script
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 
 import netlib.utils
 
@@ -32,7 +31,7 @@ class RaiseMaster(master.FlowMaster):
 def tscript(cmd, args=""):
     o = options.Options()
     cmd = example_dir.path(cmd) + " " + args
-    m = RaiseMaster(o, None, state.State())
+    m = RaiseMaster(o, None)
     sc = script.Script(cmd)
     m.addons.add(sc)
     return m, sc
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 02996881..78cfca99 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -130,7 +130,8 @@ class TestHTTPFlow:
     def test_killall(self):
         srv = DummyServer(None)
         s = flow.State()
-        fm = flow.FlowMaster(None, srv, s)
+        fm = flow.FlowMaster(None, srv)
+        fm.addons.add(s)
 
         f = tutils.tflow()
         f.reply.handle()
@@ -382,7 +383,8 @@ class TestSerialize:
     def test_load_flows(self):
         r = self._treader()
         s = flow.State()
-        fm = flow.FlowMaster(None, None, s)
+        fm = flow.FlowMaster(None, None)
+        fm.addons.add(s)
         fm.load_flows(r)
         assert len(s.flows) == 6
 
@@ -394,7 +396,8 @@ class TestSerialize:
             upstream_server="https://use-this-domain"
         )
         conf = ProxyConfig(opts)
-        fm = flow.FlowMaster(opts, DummyServer(conf), s)
+        fm = flow.FlowMaster(opts, DummyServer(conf))
+        fm.addons.add(s)
         fm.load_flows(r)
         assert s.flows[0].request.host == "use-this-domain"
 
@@ -440,8 +443,7 @@ class TestSerialize:
 class TestFlowMaster:
 
     def test_replay(self):
-        s = flow.State()
-        fm = flow.FlowMaster(None, None, s)
+        fm = flow.FlowMaster(None, None)
         f = tutils.tflow(resp=True)
         f.request.content = None
         tutils.raises("missing", fm.replay_request, f)
@@ -453,13 +455,13 @@ class TestFlowMaster:
         tutils.raises("live", fm.replay_request, f)
 
     def test_create_flow(self):
-        s = flow.State()
-        fm = flow.FlowMaster(None, None, s)
+        fm = flow.FlowMaster(None, None)
         assert fm.create_request("GET", "http", "example.com", 80, "/")
 
     def test_all(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, None, s)
+        fm = flow.FlowMaster(None, None)
+        fm.addons.add(s)
         f = tutils.tflow(req=None)
         fm.clientconnect(f.client_conn)
         f.request = HTTPRequest.wrap(netlib.tutils.treq())
diff --git a/test/mitmproxy/test_flow_state.py b/test/mitmproxy/test_flow_state.py
index cc05537d..a658e13e 100644
--- a/test/mitmproxy/test_flow_state.py
+++ b/test/mitmproxy/test_flow_state.py
@@ -5,11 +5,13 @@ from . import tutils
 class TestState:
     def test_duplicate_flow(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, None, s)
+        fm = flow.FlowMaster(None, None)
+        fm.addons.add(s)
         f = tutils.tflow(resp=True)
         fm.load_flow(f)
         assert s.flow_count() == 1
-        f2 = fm.state.duplicate_flow(f)
+
+        f2 = s.duplicate_flow(f)
         assert f2.response
         assert s.flow_count() == 2
         assert s.index(f2) == 1
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index e9db210e..0938571b 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -1,13 +1,12 @@
 import os
 import socket
 import time
-import types
 
 import netlib.tutils
 from mitmproxy import controller
 from mitmproxy import options
 from mitmproxy.builtins import script
-from mitmproxy.models import Error, HTTPResponse, HTTPFlow
+from mitmproxy.models import HTTPResponse, HTTPFlow
 from mitmproxy.proxy.config import HostMatcher, parse_server_spec
 from netlib import tcp, http, socks
 from netlib.certutils import SSLCert
@@ -986,6 +985,17 @@ class TestUpstreamProxySSL(
         assert self.chain[1].tmaster.state.flow_count() == 1
 
 
+class RequestKiller:
+    def __init__(self, exclude):
+        self.exclude = exclude
+        self.k = 0
+
+    def request(self, f):
+        self.k += 1
+        if self.k not in self.exclude:
+            f.reply.kill()
+
+
 class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest):
     ssl = True
 
@@ -995,39 +1005,18 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest):
         If we have a disconnect on a secure connection that's transparently proxified to
         an upstream http proxy, we need to send the CONNECT request again.
         """
-
-        def kill_requests(master, attr, exclude):
-            k = [0]  # variable scope workaround: put into array
-            _func = getattr(master, attr)
-
-            @controller.handler
-            def handler(*args):
-                f = args[-1]
-                k[0] += 1
-                if not (k[0] in exclude):
-                    f.client_conn.finish()
-                    f.error = Error("terminated")
-                    f.reply.kill()
-                return _func(f)
-
-            setattr(master, attr, types.MethodType(handler, master))
-
-        kill_requests(
-            self.chain[1].tmaster,
-            "request",
-            exclude = [
-                # fail first request
-                2,  # allow second request
-            ]
+        self.chain[1].tmaster.addons.add(
+            RequestKiller([2])
+        )
+        self.chain[0].tmaster.addons.add(
+            RequestKiller(
+                [
+                    1,  # CONNECT
+                    3,  # reCONNECT
+                    4   # request
+                ]
+            )
         )
-
-        kill_requests(self.chain[0].tmaster, "request",
-                      exclude=[
-                          1,  # CONNECT
-                          # fail first request
-                          3,  # reCONNECT
-                          4,  # request
-        ])
 
         p = self.pathoc()
         with p.connect():
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index 431e0f90..c4230d6f 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -7,6 +7,7 @@ import sys
 
 from mitmproxy.proxy.config import ProxyConfig
 from mitmproxy.proxy.server import ProxyServer
+from mitmproxy.flow import state
 import pathod.test
 import pathod.pathoc
 from mitmproxy import flow, controller, options
@@ -34,8 +35,9 @@ class TestMaster(flow.FlowMaster):
 
     def __init__(self, opts, config):
         s = ProxyServer(config)
-        state = flow.State()
-        flow.FlowMaster.__init__(self, opts, s, state)
+        flow.FlowMaster.__init__(self, opts, s)
+        self.state = state.State()
+        self.addons.add(self.state)
         self.addons.add(*builtins.default_addons())
         self.apps.add(testapp, "testapp", 80)
         self.apps.add(errapp, "errapp", 80)
@@ -116,8 +118,8 @@ class ProxyTestBase:
                 raise
 
     def setup(self):
-        self.master.clear_log()
         self.master.state.clear()
+        self.master.clear_log()
         self.server.clear_log()
         self.server2.clear_log()