Merge pull request #359 from MountainRider/restore-posix

Restore POSIX non_blocking() timeout support to fdpexpect.

4 files changed, 324 insertions, 36 deletions

diff --git a/pexpect/fdpexpect.py b/pexpect/fdpexpect.py
index dd1b492..40817b9 100644
--- a/pexpect/fdpexpect.py
+++ b/pexpect/fdpexpect.py
@@ -22,7 +22,8 @@ PEXPECT LICENSE
 '''
 
 from .spawnbase import SpawnBase
-from .exceptions import ExceptionPexpect
+from .exceptions import ExceptionPexpect, TIMEOUT
+from .utils import select_ignore_interrupts
 import os
 
 __all__ = ['fdspawn']
@@ -112,3 +113,27 @@ class fdspawn(SpawnBase):
         "Call self.write() for each item in sequence"
         for s in sequence:
             self.write(s)
+
+    def read_nonblocking(self, size=1, timeout=-1):
+        """ Read from the file descriptor and return the result as a string.
+
+        The read_nonblocking method of SpawnBase assumes that a call to
+        os.read will not block. This is not the case for POSIX file like
+        objects like sockets and serial ports. So we use select to implement
+        the timeout on POSIX.
+
+        :param size:    Read at most size bytes
+        :param timeout: Wait timeout seconds for file descriptor to be ready
+                        to read. If -1, use self.timeout. If 0, poll.
+        :return:        String containing the bytes read
+        """
+        if os.name == 'posix':
+            if timeout == -1:
+                timeout = self.timeout
+            rlist = [self.child_fd]
+            wlist = []
+            xlist = []
+            rlist, wlist, xlist = select_ignore_interrupts(rlist, wlist, xlist, timeout)
+            if self.child_fd not in rlist:
+                raise TIMEOUT('Timeout exceeded.')
+        return super(fdspawn, self).read_nonblocking(size)
diff --git a/pexpect/pty_spawn.py b/pexpect/pty_spawn.py
index eaf2f6f..d1c6df7 100644
--- a/pexpect/pty_spawn.py
+++ b/pexpect/pty_spawn.py
@@ -1,7 +1,6 @@
 import os
 import sys
 import time
-import select
 import pty
 import tty
 import errno
@@ -13,7 +12,7 @@ from ptyprocess.ptyprocess import use_native_pty_fork
 
 from .exceptions import ExceptionPexpect, EOF, TIMEOUT
 from .spawnbase import SpawnBase
-from .utils import which, split_command_line
+from .utils import which, split_command_line, select_ignore_interrupts
 
 @contextmanager
 def _wrap_ptyprocess_err():
@@ -403,8 +402,6 @@ class spawn(SpawnBase):
         '''
         return self.ptyproc.setecho(state)
 
-        self.echo = state
-
     def read_nonblocking(self, size=1, timeout=-1):
         '''This reads at most size characters from the child application. It
         includes a timeout. If the read does not complete within the timeout
@@ -439,7 +436,7 @@ class spawn(SpawnBase):
         # If isalive() is false, then I pretend that this is the same as EOF.
         if not self.isalive():
             # timeout of 0 means "poll"
-            r, w, e = self.__select([self.child_fd], [], [], 0)
+            r, w, e = select_ignore_interrupts([self.child_fd], [], [], 0)
             if not r:
                 self.flag_eof = True
                 raise EOF('End Of File (EOF). Braindead platform.')
@@ -447,12 +444,12 @@ class spawn(SpawnBase):
             # Irix takes a long time before it realizes a child was terminated.
             # FIXME So does this mean Irix systems are forced to always have
             # FIXME a 2 second delay when calling read_nonblocking? That sucks.
-            r, w, e = self.__select([self.child_fd], [], [], 2)
+            r, w, e = select_ignore_interrupts([self.child_fd], [], [], 2)
             if not r and not self.isalive():
                 self.flag_eof = True
                 raise EOF('End Of File (EOF). Slow platform.')
 
-        r, w, e = self.__select([self.child_fd], [], [], timeout)
+        r, w, e = select_ignore_interrupts([self.child_fd], [], [], timeout)
 
         if not r:
             if not self.isalive():
@@ -770,7 +767,7 @@ class spawn(SpawnBase):
         '''
 
         while self.isalive():
-            r, w, e = self.__select([self.child_fd, self.STDIN_FILENO], [], [])
+            r, w, e = select_ignore_interrupts([self.child_fd, self.STDIN_FILENO], [], [])
             if self.child_fd in r:
                 try:
                     data = self.__interact_read(self.child_fd)
@@ -802,33 +799,6 @@ class spawn(SpawnBase):
                 self._log(data, 'send')
                 self.__interact_writen(self.child_fd, data)
 
-    def __select(self, iwtd, owtd, ewtd, timeout=None):
-
-        '''This is a wrapper around select.select() that ignores signals. If
-        select.select raises a select.error exception and errno is an EINTR
-        error then it is ignored. Mainly this is used to ignore sigwinch
-        (terminal resize). '''
-
-        # if select() is interrupted by a signal (errno==EINTR) then
-        # we loop back and enter the select() again.
-        if timeout is not None:
-            end_time = time.time() + timeout
-        while True:
-            try:
-                return select.select(iwtd, owtd, ewtd, timeout)
-            except select.error:
-                err = sys.exc_info()[1]
-                if err.args[0] == errno.EINTR:
-                    # if we loop back we have to subtract the
-                    # amount of time we already waited.
-                    if timeout is not None:
-                        timeout = end_time - time.time()
-                        if timeout < 0:
-                            return([], [], [])
-                else:
-                    # something else caused the select.error, so
-                    # this actually is an exception.
-                    raise
 
 def spawnu(*args, **kwargs):
     """Deprecated: pass encoding to spawn() instead."""
diff --git a/pexpect/utils.py b/pexpect/utils.py
index c2763c4..ae0fe9d 100644
--- a/pexpect/utils.py
+++ b/pexpect/utils.py
@@ -1,6 +1,15 @@
 import os
 import sys
 import stat
+import select
+import time
+import errno
+
+try:
+    InterruptedError
+except NameError:
+    # Alias Python2 exception to Python3
+    InterruptedError = select.error
 
 
 def is_executable_file(path):
@@ -111,3 +120,32 @@ def split_command_line(command_line):
     if arg != '':
         arg_list.append(arg)
     return arg_list
+
+
+def select_ignore_interrupts(iwtd, owtd, ewtd, timeout=None):
+
+    '''This is a wrapper around select.select() that ignores signals. If
+    select.select raises a select.error exception and errno is an EINTR
+    error then it is ignored. Mainly this is used to ignore sigwinch
+    (terminal resize). '''
+
+    # if select() is interrupted by a signal (errno==EINTR) then
+    # we loop back and enter the select() again.
+    if timeout is not None:
+        end_time = time.time() + timeout
+    while True:
+        try:
+            return select.select(iwtd, owtd, ewtd, timeout)
+        except InterruptedError:
+            err = sys.exc_info()[1]
+            if err.args[0] == errno.EINTR:
+                # if we loop back we have to subtract the
+                # amount of time we already waited.
+                if timeout is not None:
+                    timeout = end_time - time.time()
+                    if timeout < 0:
+                        return([], [], [])
+            else:
+                # something else caused the select.error, so
+                # this actually is an exception.
+                raise
diff --git a/tests/test_socket.py b/tests/test_socket.py
new file mode 100644
index 0000000..56d69c7
--- /dev/null
+++ b/tests/test_socket.py
@@ -0,0 +1,255 @@
+#!/usr/bin/env python
+'''
+PEXPECT LICENSE
+
+    This license is approved by the OSI and FSF as GPL-compatible.
+        http://opensource.org/licenses/isc-license.txt
+
+    Copyright (c) 2012, Noah Spurrier <noah@noah.org>
+    PERMISSION TO USE, COPY, MODIFY, AND/OR DISTRIBUTE THIS SOFTWARE FOR ANY
+    PURPOSE WITH OR WITHOUT FEE IS HEREBY GRANTED, PROVIDED THAT THE ABOVE
+    COPYRIGHT NOTICE AND THIS PERMISSION NOTICE APPEAR IN ALL COPIES.
+    THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+    WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+    MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+    ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+    WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+    ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+    OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+'''
+import pexpect
+from pexpect import fdpexpect
+import unittest
+from . import PexpectTestCase
+import multiprocessing
+import os
+import signal
+import socket
+import time
+import errno
+
+
+class SocketServerError(Exception):
+    pass
+
+
+class ExpectTestCase(PexpectTestCase.PexpectTestCase):
+
+    def setUp(self):
+        print(self.id())
+        PexpectTestCase.PexpectTestCase.setUp(self)
+        self.host = '127.0.0.1'
+        self.port = 49152 + 10000
+        self.motd = b"""\
+------------------------------------------------------------------------------
+*                  Welcome to the SOCKET UNIT TEST code!                     *
+------------------------------------------------------------------------------
+*                                                                            *
+* This unit test code is our best effort at testing the ability of the       *
+* pexpect library to handle sockets. We need some text to test buffer size   *
+* handling.                                                                  *
+*                                                                            *
+* A page is 1024 bytes or 1K. 80 x 24 = 1920. So a standard terminal window  *
+* contains more than one page. We actually want more than a page for our     *
+* tests.                                                                     *
+*                                                                            *
+* This is the twelfth line, and we need 24. So we need a few more paragraphs.*
+* We can keep them short and just put lines between them.                    *
+*                                                                            *
+* The 80 x 24 terminal size comes from the ancient past when computers were  *
+* only able to display text in cuneiform writing.                            *
+*                                                                            *
+* The cunieform writing system used the edge of a reed to make marks on clay *
+* tablets.                                                                   *
+*                                                                            *
+* It was the forerunner of the style of handwriting used by doctors to write *
+* prescriptions. Thus the name: pre (before) script (writing) ion (charged   *
+* particle).                                                                 *
+------------------------------------------------------------------------------
+""".replace(b'\n', b'\n\r') + b"\r\n"
+        self.prompt1 = b'Press Return to continue:'
+        self.prompt2 = b'Rate this unit test>'
+        self.prompt3 = b'Press X to exit:'
+        self.enter = b'\r\n'
+        self.exit = b'X\r\n'
+        self.server_up = multiprocessing.Event()
+        self.server_process = multiprocessing.Process(target=self.socket_server, args=(self.server_up,))
+        self.server_process.daemon = True
+        self.server_process.start()
+        counter = 0
+        while not self.server_up.is_set():
+            time.sleep(0.250)
+            counter += 1
+            if counter > (10 / 0.250):
+                raise SocketServerError("Could not start socket server")
+
+    def tearDown(self):
+        os.kill(self.server_process.pid, signal.SIGINT)
+        self.server_process.join(timeout=5.0)
+        PexpectTestCase.PexpectTestCase.tearDown(self)
+
+    def socket_server(self, server_up):
+        sock = None
+        try:
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            sock.bind((self.host, self.port))
+            sock.listen(5)
+            server_up.set()
+            while True:
+                (conn, addr) = sock.accept()
+                conn.send(self.motd)
+                conn.send(self.prompt1)
+                result = conn.recv(1024)
+                if result != self.enter:
+                    break
+                conn.send(self.prompt2)
+                result = conn.recv(1024)
+                if result != self.enter:
+                    break
+                conn.send(self.prompt3)
+                result = conn.recv(1024)
+                if result.startswith(self.exit[0]):
+                    conn.shutdown(socket.SHUT_RDWR)
+                    conn.close()
+        except KeyboardInterrupt:
+            pass
+        if sock is not None:
+            try:
+                sock.shutdown(socket.SHUT_RDWR)
+                sock.close()
+            except socket.error:
+                pass
+        exit(0)
+
+    def socket_fn(self, timed_out, all_read):
+        result = 0
+        try:
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.connect((self.host, self.port))
+            session = fdpexpect.fdspawn(sock, timeout=10)
+            # Get all data from server
+            session.read_nonblocking(size=4096)
+            all_read.set()
+            # This read should timeout
+            session.read_nonblocking(size=4096)
+        except pexpect.TIMEOUT:
+            timed_out.set()
+            result = errno.ETIMEDOUT
+        exit(result)
+
+    def test_socket(self):
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.connect((self.host, self.port))
+        session = fdpexpect.fdspawn(sock.fileno(), timeout=10)
+        session.expect(self.prompt1)
+        self.assertEqual(session.before, self.motd)
+        session.send(self.enter)
+        session.expect(self.prompt2)
+        session.send(self.enter)
+        session.expect(self.prompt3)
+        session.send(self.exit)
+        session.expect(pexpect.EOF)
+        self.assertEqual(session.before, b'')
+
+    def test_socket_with_write(self):
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.connect((self.host, self.port))
+        session = fdpexpect.fdspawn(sock.fileno(), timeout=10)
+        session.expect(self.prompt1)
+        self.assertEqual(session.before, self.motd)
+        session.write(self.enter)
+        session.expect(self.prompt2)
+        session.write(self.enter)
+        session.expect(self.prompt3)
+        session.write(self.exit)
+        session.expect(pexpect.EOF)
+        self.assertEqual(session.before, b'')
+
+    def test_not_int(self):
+        with self.assertRaises(pexpect.ExceptionPexpect):
+            session = fdpexpect.fdspawn('bogus', timeout=10)
+
+    def test_not_file_descriptor(self):
+        with self.assertRaises(pexpect.ExceptionPexpect):
+            session = fdpexpect.fdspawn(-1, timeout=10)
+
+    def test_timeout(self):
+        with self.assertRaises(pexpect.TIMEOUT):
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.connect((self.host, self.port))
+            session = fdpexpect.fdspawn(sock, timeout=10)
+            session.expect(b'Bogus response')
+
+    def test_interrupt(self):
+        timed_out = multiprocessing.Event()
+        all_read = multiprocessing.Event()
+        test_proc = multiprocessing.Process(target=self.socket_fn, args=(timed_out, all_read))
+        test_proc.daemon = True
+        test_proc.start()
+        while not all_read.is_set():
+            time.sleep(1.0)
+        os.kill(test_proc.pid, signal.SIGWINCH)
+        while not timed_out.is_set():
+            time.sleep(1.0)
+        test_proc.join(timeout=5.0)
+        self.assertEqual(test_proc.exitcode, errno.ETIMEDOUT)
+
+    def test_multiple_interrupts(self):
+        timed_out = multiprocessing.Event()
+        all_read = multiprocessing.Event()
+        test_proc = multiprocessing.Process(target=self.socket_fn, args=(timed_out, all_read))
+        test_proc.daemon = True
+        test_proc.start()
+        while not all_read.is_set():
+            time.sleep(1.0)
+        while not timed_out.is_set():
+            os.kill(test_proc.pid, signal.SIGWINCH)
+            time.sleep(1.0)
+        test_proc.join(timeout=5.0)
+        self.assertEqual(test_proc.exitcode, errno.ETIMEDOUT)
+
+    def test_maxread(self):
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.connect((self.host, self.port))
+        session = fdpexpect.fdspawn(sock.fileno(), timeout=10)
+        session.maxread = 1100
+        session.expect(self.prompt1)
+        self.assertEqual(session.before, self.motd)
+        session.send(self.enter)
+        session.expect(self.prompt2)
+        session.send(self.enter)
+        session.expect(self.prompt3)
+        session.send(self.exit)
+        session.expect(pexpect.EOF)
+        self.assertEqual(session.before, b'')
+
+    def test_fd_isalive (self):
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.connect((self.host, self.port))
+        session = fdpexpect.fdspawn(sock.fileno(), timeout=10)
+        assert session.isalive()
+        sock.close()
+        assert not session.isalive(), "Should not be alive after close()"
+
+    def test_fd_isatty (self):
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.connect((self.host, self.port))
+        session = fdpexpect.fdspawn(sock.fileno(), timeout=10)
+        assert not session.isatty()
+        session.close()
+
+    def test_fileobj(self):
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.connect((self.host, self.port))
+        session = fdpexpect.fdspawn(sock, timeout=10) # Should get the fileno from the socket
+        session.expect(self.prompt1)
+        session.close()
+        assert not session.isalive()
+        session.close()  # Smoketest - should be able to call this again
+
+if __name__ == '__main__':
+    unittest.main()
+
+suite = unittest.makeSuite(ExpectTestCase, 'test')