diff --git a/netlib/tcp.py b/netlib/tcp.py index 57a9b737f..1523370bd 100644 --- a/netlib/tcp.py +++ b/netlib/tcp.py @@ -272,7 +272,7 @@ class Reader(_FileLike): Raises: TcpException if there was an error with the socket TlsException if there was an error with pyOpenSSL. - NotImplementedError if the underlying file object is not a (pyOpenSSL) socket + NotImplementedError if the underlying file object is not a [pyOpenSSL] socket """ if isinstance(self.o, socket_fileobject): try: diff --git a/test/test_tcp.py b/test/test_tcp.py index 20a295ddd..2b091ef01 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): @@ -716,15 +716,34 @@ class TestFileLike: 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)) - c.connect() + self._connect(c) c.wfile.write(testval) c.wfile.flush() - assert c.rfile.peek(4) == b"peek"[:4] - assert c.rfile.peek(6) == testval + 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: