Refator server tests to use flow.FlowMaster and flow.State

3 files changed, 33 insertions, 43 deletions

diff --git a/test/test_proxy.py b/test/test_proxy.py
index 08b3634f..d4d275b5 100644
--- a/test/test_proxy.py
+++ b/test/test_proxy.py
@@ -5,21 +5,19 @@ from libmproxy import proxy, flow
 import tutils
 
 
-class TestProxyError:
-    def test_simple(self):
-        p = proxy.ProxyError(111, "msg")
-        assert repr(p)
+def test_proxy_error():
+    p = proxy.ProxyError(111, "msg")
+    assert str(p)
 
 
-class TestAppRegistry:
-    def test_add_get(self):
-        ar = proxy.AppRegistry()
-        ar.add("foo", "domain", 80)
+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 = tutils.treq()
+    r.host = "domain"
+    r.port = 80
+    assert ar.get(r)
 
-        r.port = 81
-        assert not ar.get(r)
+    r.port = 81
+    assert not ar.get(r)
diff --git a/test/test_server.py b/test/test_server.py
index 1f83c496..87996c43 100644
--- a/test/test_server.py
+++ b/test/test_server.py
@@ -15,7 +15,7 @@ import tutils
 class SanityMixin:
     def test_http(self):
         assert self.pathod("304").status_code == 304
-        assert self.log()
+        assert self.master.state.view
 
     def test_large(self):
         assert len(self.pathod("200:b@50k").content) == 1024*50
@@ -54,7 +54,7 @@ class TestProxy(tutils.HTTPProxTest):
         f = self.pathod("304")
         assert f.status_code == 304
 
-        l = self.log()
-        assert l[1].address
-        assert "host" in l[2].headers
-        assert l[3].code == 304
+        l = self.master.state.view[0]
+        assert l.request.client_conn.address
+        assert "host" in l.request.headers
+        assert l.response.code == 304
diff --git a/test/tutils.py b/test/tutils.py
index 79a1da39..a35dc4d4 100644
--- a/test/tutils.py
+++ b/test/tutils.py
@@ -43,28 +43,28 @@ def tflow_err():
     return f
 
 
-class TestMaster(controller.Master):
-    def __init__(self, port, testq, config):
-        s = proxy.ProxyServer(config, port)
-        controller.Master.__init__(self, s)
+class TestMaster(flow.FlowMaster):
+    def __init__(self, testq, config):
+        s = proxy.ProxyServer(config, 0)
+        state = flow.State()
+        flow.FlowMaster.__init__(self, s, state)
         self.testq = testq
-        self.log = []
-
-    def clear(self):
-        self.log = []
 
     def handle(self, m):
-        self.log.append(m)
+        f = flow.FlowMaster.handle(self, m)
         m._ack()
 
 
 class ProxyThread(threading.Thread):
     def __init__(self, testq, config):
-        self.port = random.randint(10000, 20000)
-        self.tmaster = TestMaster(self.port, testq, config)
+        self.tmaster = TestMaster(testq, config)
         controller.should_exit = False
         threading.Thread.__init__(self)
 
+    @property
+    def port(self):
+        return self.tmaster.server.port
+
     def run(self):
         self.tmaster.run()
 
@@ -72,18 +72,6 @@ class ProxyThread(threading.Thread):
         self.tmaster.shutdown()
 
 
-class ServerThread(threading.Thread):
-    def __init__(self, server):
-        self.server = server
-        threading.Thread.__init__(self)
-
-    def run(self):
-        self.server.serve_forever()
-
-    def shutdown(self):
-        self.server.shutdown()
-
-
 class ProxTestBase:
     @classmethod
     def setupAll(cls):
@@ -97,13 +85,17 @@ class ProxTestBase:
         cls.proxy = ProxyThread(cls.tqueue, config)
         cls.proxy.start()
 
+    @property
+    def master(cls):
+        return cls.proxy.tmaster
+
     @classmethod
     def teardownAll(cls):
         cls.proxy.shutdown()
         cls.server.shutdown()
 
     def setUp(self):
-        self.proxy.tmaster.clear()
+        self.master.state.clear()
 
     @property
     def scheme(self):