Test request and response kill functionality.

5 files changed, 40 insertions, 11 deletions

diff --git a/.coveragerc b/.coveragerc
index 696e0eb8..dd57a6e7 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -1,3 +1,6 @@
+[rum]
+branch = True
+
 [report]
 omit = *contrib*, *tnetstring*, *platform*
 include = *libmproxy*
diff --git a/libmproxy/controller.py b/libmproxy/controller.py
index da097692..bb22597d 100644
--- a/libmproxy/controller.py
+++ b/libmproxy/controller.py
@@ -41,10 +41,13 @@ class Reply:
         self.q = Queue.Queue()
         self.acked = False
 
-    def __call__(self, msg=False):
+    def __call__(self, msg=None):
         if not self.acked:
             self.acked = True
-            self.q.put(msg or self.obj)
+            if msg is None:
+                self.q.put(self.obj)
+            else:
+                self.q.put(msg)
 
 
 class Channel:
@@ -62,7 +65,7 @@ class Channel:
             try:
                 # The timeout is here so we can handle a should_exit event.
                 g = m.reply.q.get(timeout=0.5)
-            except Queue.Empty:
+            except Queue.Empty: # pragma: nocover
                 continue
             return g
 
diff --git a/test/.gitignore b/test/.gitignore
deleted file mode 100644
index 6350e986..00000000
--- a/test/.gitignore
+++ /dev/null
@@ -1 +0,0 @@
-.coverage
diff --git a/test/.pry b/test/.pry
deleted file mode 100644
index f6f18e7b..00000000
--- a/test/.pry
+++ /dev/null
@@ -1,6 +0,0 @@
-base = ..
-coverage = ../libmproxy
-exclude = .
-          ../libmproxy/contrib
-          ../libmproxy/tnetstring.py
-
diff --git a/test/test_server.py b/test/test_server.py
index 5cba891c..8aefa4b8 100644
--- a/test/test_server.py
+++ b/test/test_server.py
@@ -2,7 +2,7 @@ import socket, time
 from netlib import tcp
 from libpathod import pathoc
 import tutils, tservers
-from libmproxy import flow
+from libmproxy import flow, proxy
 
 """
     Note that the choice of response code in these tests matters more than you
@@ -147,6 +147,7 @@ class TestProxy(tservers.HTTPProxTest):
         assert request.timestamp_end - request.timestamp_start <= 0.1
 
 
+
 class MasterFakeResponse(tservers.TestMaster):
     def handle_request(self, m):
         resp = tutils.tresp()
@@ -160,3 +161,32 @@ class TestFakeResponse(tservers.HTTPProxTest):
         f = self.pathod("200")
         assert "header_response" in f.headers.keys()
 
+
+
+class MasterKillRequest(tservers.TestMaster):
+    def handle_request(self, m):
+        m.reply(proxy.KILL)
+
+
+class TestKillRequest(tservers.HTTPProxTest):
+    masterclass = MasterKillRequest
+    def test_kill(self):
+        p = self.pathoc()
+        tutils.raises("empty reply", self.pathod, "200")
+        # Nothing should have hit the server
+        assert not self.last_log()
+
+
+class MasterKillResponse(tservers.TestMaster):
+    def handle_response(self, m):
+        m.reply(proxy.KILL)
+
+
+class TestKillResponse(tservers.HTTPProxTest):
+    masterclass = MasterKillResponse
+    def test_kill(self):
+        p = self.pathoc()
+        tutils.raises("empty reply", self.pathod, "200")
+        # The server should have seen a request
+        assert self.last_log()
+