Unit test love - 100% for flow.py, dump.py

libmproxy/flow.py

@@ -1375,6 +1375,8 @@ class FlowMaster(controller.Master):
         self.kill_nonreplay = kill
 
     def stop_server_playback(self):
+        if self.server_playback.exit:
+            self.shutdown()
         self.server_playback = None
 
     def do_server_playback(self, flow):
@@ -1408,10 +1410,6 @@ class FlowMaster(controller.Master):
                 self.shutdown()
             self.client_playback.tick(self)
 
-        if self.server_playback:
-            if self.server_playback.exit and self.server_playback.count() == 0:
-                self.shutdown()
-
         return controller.Master.tick(self, q)
 
     def duplicate_flow(self, f):

test/test_dump.py

@@ -48,6 +48,15 @@ class TestDumpMaster:
         fw.add(t)
         f.close()
 
+    def test_error(self):
+        cs = StringIO()
+        o = dump.Options(verbosity=1)
+        m = dump.DumpMaster(None, o, None, outfile=cs)
+        f = tutils.tflow_err()
+        m.handle_request(f.request)
+        assert m.handle_error(f.error)
+        assert "error" in cs.getvalue()
+
     def test_replay(self):
         cs = StringIO()
 

test/test_flow.py

@@ -624,6 +624,7 @@ class TestFlowMaster:
         fm.handle_error(flow.Error(f.request, "error"))
 
     def test_server_playback(self):
+        controller.should_exit = False
         s = flow.State()
 
         f = tutils.tflow()
@@ -641,8 +642,9 @@ class TestFlowMaster:
         r = tutils.tflow()
         r.request.content = "gibble"
         assert not fm.do_server_playback(r)
-
         assert fm.do_server_playback(tutils.tflow())
+
+        fm.start_server_playback(pb, False, [], True, False)
         q = Queue.Queue()
         fm.tick(q)
         assert controller.should_exit
@@ -650,6 +652,20 @@ class TestFlowMaster:
         fm.stop_server_playback()
         assert not fm.server_playback
 
+    def test_server_playback_kill(self):
+        s = flow.State()
+        f = tutils.tflow()
+        f.response = tutils.tresp(f.request)
+        pb = [f]
+        fm = flow.FlowMaster(None, s)
+        fm.refresh_server_playback = True
+        fm.start_server_playback(pb, True, [], False, False)
+
+        f = tutils.tflow()
+        f.request.host = "nonexistent"
+        fm.process_new_request(f)
+        assert "killed" in f.error.msg
+
     def test_stickycookie(self):
         s = flow.State()
         fm = flow.FlowMaster(None, s)