100% test coverage

2 files changed, 22 insertions, 22 deletions

diff --git a/libpathod/pathoc.py b/libpathod/pathoc.py
index ca21a423..4efa0447 100644
--- a/libpathod/pathoc.py
+++ b/libpathod/pathoc.py
@@ -293,35 +293,19 @@ class Pathoc(tcp.TCPClient):
                     return
                 yield frm
 
-    def websocket_get_frame(self, frame):
-        """
-            Called when a frame is received from the server.
-        """
-        pass
-
     def websocket_send_frame(self, r):
         """
             Sends a single websocket frame.
         """
         with self.log() as log:
-            if isinstance(r, basestring):
-                r = language.parse_pathoc(r).next()
             log(">> %s" % r)
-            try:
-                language.serve(r, self.wfile, self.settings)
-                self.wfile.flush()
-            except tcp.NetLibTimeout:
-                if self.ignoretimeout:
-                    self.log("Timeout (ignored)")
-                    return None
-                raise
+            language.serve(r, self.wfile, self.settings)
+            self.wfile.flush()
 
-    def websocket_start(self, r, limit=None):
+    def websocket_start(self, r):
         """
             Performs an HTTP request, and attempts to drop into websocket
             connection.
-
-            limit: Disconnect after receiving N server frames.
         """
         resp = self.http(r)
         if resp.status_code == 101:
@@ -348,8 +332,6 @@ class Pathoc(tcp.TCPClient):
             May raise http.HTTPError, tcp.NetLibError
         """
         with self.log() as log:
-            if isinstance(r, basestring):
-                r = language.parse_pathoc(r).next()
             log(">> %s" % r)
             resp, req = None, None
             try:
@@ -395,7 +377,7 @@ class Pathoc(tcp.TCPClient):
             r = language.parse_pathoc(r).next()
         if isinstance(r, language.http.Request):
             if r.ws:
-                return self.websocket_start(r, self.websocket_get_frame)
+                return self.websocket_start(r)
             else:
                 return self.http(r)
         elif isinstance(r, language.websockets.WebsocketFrame):
diff --git a/test/test_pathoc.py b/test/test_pathoc.py
index 21e46e8c..d37cf9eb 100644
--- a/test/test_pathoc.py
+++ b/test/test_pathoc.py
@@ -185,6 +185,24 @@ class TestDaemon(_TestDaemon):
     def test_conn_err(self):
         assert "Invalid server response" in self.tval(["get:'/p/200:d2'"])
 
+    def test_websocket_shutdown(self):
+        c = pathoc.Pathoc(("127.0.0.1", self.d.port), fp=None)
+        c.connect()
+        c.request("ws:/")
+        c.stop()
+
+    def test_wait_finish(self):
+        c = pathoc.Pathoc(
+            ("127.0.0.1", self.d.port),
+            fp=None,
+            ws_read_limit = 1
+        )
+        c.connect()
+        c.request("ws:/")
+        c.request("wf:f'wf:x100'")
+        [i for i in c.wait(timeout=0, finish=False)]
+        [i for i in c.wait(timeout=0)]
+
     def test_connect_fail(self):
         to = ("foobar", 80)
         c = pathoc.Pathoc(("127.0.0.1", self.d.port), fp=None)