5 files changed, 25 insertions, 24 deletions
diff --git a/test/test_controller.py b/test/test_controller.py
index f6d6b5eb..e71a148e 100644
--- a/test/test_controller.py
+++ b/test/test_controller.py
@@ -6,7 +6,7 @@ class TestMaster:
     def test_default_handler(self):
         m = controller.Master(None)
         msg = mock.MagicMock()
-        m.handle(msg)
+        m.handle("type", msg)
         assert msg.reply.call_count == 1
 
 
diff --git a/test/test_dump.py b/test/test_dump.py
index 3d375f16..9874b650 100644
--- a/test/test_dump.py
+++ b/test/test_dump.py
@@ -114,7 +114,7 @@ class TestDumpMaster:
         o = dump.Options(app=True)
         s = mock.MagicMock()
         m = dump.DumpMaster(s, o, None)
-        assert s.apps.add.call_count == 1
+        assert len(m.apps.apps) == 1
 
     def test_replacements(self):
         o = dump.Options(replacements=[(".*", "content", "foo")])
diff --git a/test/test_flow.py b/test/test_flow.py
index c614960b..bf6a7a42 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -5,6 +5,27 @@ from libmproxy import filt, flow, controller, utils, tnetstring, proxy
 import tutils
 
 
+def test_app_registry():
+    ar = flow.AppRegistry()
+    ar.add("foo", "domain", 80)
+
+    r = tutils.treq()
+    r.host = "domain"
+    r.port = 80
+    assert ar.get(r)
+
+    r.port = 81
+    assert not ar.get(r)
+
+    r = tutils.treq()
+    r.host = "domain2"
+    r.port = 80
+    assert not ar.get(r)
+    r.headers["host"] = ["domain"]
+    assert ar.get(r)
+
+
+
 class TestStickyCookieState:
     def _response(self, cookie, host):
         s = flow.StickyCookieState(filt.parse(".*"))
diff --git a/test/test_proxy.py b/test/test_proxy.py
index dfccaaf7..371e5ef7 100644
--- a/test/test_proxy.py
+++ b/test/test_proxy.py
@@ -11,26 +11,6 @@ def test_proxy_error():
     assert str(p)
 
 
-def test_app_registry():
-    ar = proxy.AppRegistry()
-    ar.add("foo", "domain", 80)
-
-    r = tutils.treq()
-    r.host = "domain"
-    r.port = 80
-    assert ar.get(r)
-
-    r.port = 81
-    assert not ar.get(r)
-
-    r = tutils.treq()
-    r.host = "domain2"
-    r.port = 80
-    assert not ar.get(r)
-    r.headers["host"] = ["domain"]
-    assert ar.get(r)
-
-
 class TestServerConnection:
     def setUp(self):
         self.d = test.Daemon()
diff --git a/test/tservers.py b/test/tservers.py
index f886d42d..ac95b168 100644
--- a/test/tservers.py
+++ b/test/tservers.py
@@ -23,10 +23,10 @@ def errapp(environ, start_response):
 class TestMaster(flow.FlowMaster):
     def __init__(self, testq, config):
         s = proxy.ProxyServer(config, 0)
-        s.apps.add(testapp, "testapp", 80)
-        s.apps.add(errapp, "errapp", 80)
         state = flow.State()
         flow.FlowMaster.__init__(self, s, state)
+        self.apps.add(testapp, "testapp", 80)
+        self.apps.add(errapp, "errapp", 80)
         self.testq = testq
         self.clear_log()
         self.start_app(APP_HOST, APP_PORT, False)