Merge pull request #118 from mitmproxy/py3-peek

Fix Reader.peek() on Python 3

1 files changed, 34 insertions, 1 deletions

diff --git a/test/test_tcp.py b/test/test_tcp.py
index 738fb2eb..2b091ef0 100644
--- a/test/test_tcp.py
+++ b/test/test_tcp.py
@@ -12,7 +12,7 @@ import OpenSSL
 
 from netlib import tcp, certutils, tutils, tservers
 from netlib.exceptions import InvalidCertificateException, TcpReadIncomplete, TlsException, \
-    TcpTimeout, TcpDisconnect, TcpException
+    TcpTimeout, TcpDisconnect, TcpException, NetlibException
 
 
 class EchoHandler(tcp.BaseHandler):
@@ -713,6 +713,39 @@ class TestFileLike:
         tutils.raises(TcpReadIncomplete, s.safe_read, 10)
 
 
+class TestPeek(tservers.ServerTestBase):
+    handler = EchoHandler
+
+    def _connect(self, c):
+        c.connect()
+
+    def test_peek(self):
+        testval = b"peek!\n"
+        c = tcp.TCPClient(("127.0.0.1", self.port))
+        self._connect(c)
+        c.wfile.write(testval)
+        c.wfile.flush()
+
+        assert c.rfile.peek(4) == b"peek"
+        assert c.rfile.peek(6) == b"peek!\n"
+        assert c.rfile.readline() == testval
+
+        c.close()
+        with tutils.raises(NetlibException):
+            if c.rfile.peek(1) == b"":
+                # Workaround for Python 2 on Unix:
+                # Peeking a closed connection does not raise an exception here.
+                raise NetlibException()
+
+
+class TestPeekSSL(TestPeek):
+    ssl = True
+
+    def _connect(self, c):
+        c.connect()
+        c.convert_to_ssl()
+
+
 class TestAddress:
 
     def test_simple(self):