first_stage_test: Add more tests test_stdin_non_blocking and test_stdin_blocking

While at it, refactor the code, convert the inline comments to docstrings. Those
docstrings are later reused as the test description by the test runner.

Signed-off-by: Marc Hartmayer <[email protected]>

Author: mhartmay

tests/first_stage_test.py

@@ -1,3 +1,6 @@
+import fcntl
+import os
+import select
 import subprocess
 import sys
 
@@ -17,27 +20,36 @@ class CommandLineTest(testlib.RouterMixin, testlib.TestCase):
     #   * 2.7 starting 3.x
     #   * 3.x starting 2.7
 
-    def test_valid_syntax(self):
+    def setUp(self):
+        super(CommandLineTest, self).setUp()
         options = mitogen.parent.Options(max_message_size=123)
         conn = mitogen.parent.Connection(options, self.router)
         conn.context = mitogen.core.Context(None, 123)
-        args = conn.get_boot_command()
+        self.args = conn.get_boot_command()
+        self.preamble = conn.get_preamble()
+        self.conn = conn
+
+    def test_valid_syntax(self):
+        """Test valid syntax
 
-        # The boot command should write an ECO marker to stdout, read the
-        # preamble from stdin, then execute it.
+        The boot command should write an ECO marker to stdout, read the
+        preamble from stdin, then execute it.
 
-        # This test attaches /dev/zero to stdin to create a specific failure
-        # 1. Fork child reads <compressed preamble size> bytes of NUL (`b'\0'`)
-        # 2. Fork child crashes (trying to decompress the junk data)
-        # 3. Fork child's file descriptors (write pipes) are closed by the OS
-        # 4. Fork parent does `dup(<read pipe>, <stdin>)` and `exec(<python>)`
-        # 5. Python reads `b''` (i.e. EOF) from stdin (a closed pipe)
-        # 6. Python runs `''` (a valid script) and exits with success
+        This test attaches /dev/zero to stdin to create a specific failure
+
+        1. Fork child reads <compressed preamble size> bytes of NUL (`b'\0'`)
+        2. Fork child crashes (trying to decompress the junk data)
+        3. Fork child's file descriptors (write pipes) are closed by the OS
+        4. Fork parent does `dup(<read pipe>, <stdin>)` and `exec(<python>)`
+        5. Python reads `b''` (i.e. EOF) from stdin (a closed pipe)
+        6. Python runs `''` (a valid script) and exits with success
+
+        """
 
         fp = open("/dev/zero", "r")
         try:
             proc = subprocess.Popen(
-                args,
+                self.args,
                 stdin=fp,
                 stdout=subprocess.PIPE,
                 stderr=subprocess.PIPE,
@@ -55,32 +67,29 @@ def test_valid_syntax(self):
             fp.close()
 
     def test_stage(self):
-        options = mitogen.parent.Options(max_message_size=123)
-        conn = mitogen.parent.Connection(options, self.router)
-        conn.context = mitogen.core.Context(None, 123)
-        args = conn.get_boot_command()
-        preamble = conn.get_preamble()
+        """Test that first stage works
+
+        The boot command should read the preamble from STDIN, write all ECO
+        markers to STDOUT, and then execute the preamble.
 
-        # The boot command should write all ECO markers to stdout, read the
-        # preamble from stdin, then execute it.
+        This test writes the complete preamble to STDIN.
 
-        # This test writes the preamble to STDIN and closes it then to create an
-        # EOF situation.
-        # 1. Fork child tries to read from STDIN, but stops as EOF is received.
-        # 2. Fork child writes all ECO markers to stdout
-        # TBD
+        1. Fork child reads from STDIN
+        2. Fork child writes all ECO markers to stdout as expected.
+
+        """
 
         proc = subprocess.Popen(
-            args=args,
+            args=self.args,
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             stdin=subprocess.PIPE,
         )
         try:
             try:
-                stdout, stderr = proc.communicate(input=preamble, timeout=10)
+                stdout, stderr = proc.communicate(input=self.preamble, timeout=10)
             except TypeError:
-                stdout, stderr = proc.communicate(input=preamble)
+                stdout, stderr = proc.communicate(input=self.preamble)
         except:
             proc.kill()
             self.fail("First stage did not finish")
@@ -101,34 +110,172 @@ def test_stage(self):
                 stderr,
             )
 
-    def test_eof_too_early(self):
-        options = mitogen.parent.Options(max_message_size=123)
-        conn = mitogen.parent.Connection(options, self.router)
-        conn.context = mitogen.core.Context(None, 123)
-        args = conn.get_boot_command()
-        preamble = conn.get_preamble()
+    def test_stdin_non_blocking(self):
+        """Test that first stage works with non-blocking STDIN
+
+        The boot command should read the preamble from STDIN, write all ECO
+        markers to STDOUT, and then execute the preamble.
+
+        This test writes the complete preamble to non-blocking STDIN.
 
-        # The boot command should write an ECO marker to stdout, read the
-        # preamble from stdin, then execute it.
+        1. Fork child reads from non-blocking STDIN
+        2. Fork child writes all ECO markers to stdout as expected.
 
-        # This test writes some data to STDIN and closes it then to create an
-        # EOF situation.
-        # 1. Fork child tries to read from STDIN, but stops as EOF is received.
-        # 2. Fork child crashes (trying to decompress the junk data)
-        # 3. Fork child's file descriptors (write pipes) are closed by the OS
-        # 4. Fork parent does `dup(<read pipe>, <stdin>)` and `exec(<python>)`
-        # 5. Python reads `b''` (i.e. EOF) from stdin (a closed pipe)
-        # 6. Python runs `''` (a valid script) and exits with success
+        """
+
+        PIPE_SIZE = 4096
+        # Make sure that not all data can be written in one write operation.
+        CHUNK_SIZE = 2 * PIPE_SIZE
+        r, w = mitogen.core.pipe(blocking=False)
+        with w:
+            try:
+                fcntl.fcntl(w.fileno(), fcntl.F_SETPIPE_SZ, PIPE_SIZE)
+            except AttributeError:
+                pass
+            else:
+                self.assertEqual(fcntl.fcntl(w.fileno(), fcntl.F_GETPIPE_SZ), PIPE_SIZE)
+
+            proc = subprocess.Popen(
+                args=self.args,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                stdin=r,
+                close_fds=True,
+            )
+            # Close the read fd in the parent
+            r.close()
+
+            view = mitogen.core.BufferType(self.preamble, 0)
+            offset = 0
+            while offset < len(view):
+                end = min(offset + CHUNK_SIZE, len(view))
+                _, w_fds, _ = select.select([], [w.fileno()], [], 10)
+                self.assertIn(w.fileno(), w_fds)
+                written = os.write(w.fileno(), view[offset:end])
+                self.assertGreater(written, 0)
+                offset += written
+
+        try:
+            try:
+                returncode = proc.wait(timeout=10)
+            except TypeError:
+                returncode = proc.wait()
+        except:
+            proc.kill()
+            self.fail("First stage did not finish")
+        else:
+            try:
+                # proc was killed by SIGTERM?!? TODO Where does this come from?
+                self.assertEqual(-15, proc.returncode)
+                self.assertEqual(
+                    proc.stdout.read(),
+                    mitogen.parent.BootstrapProtocol.EC0_MARKER
+                    + b("\n")
+                    + mitogen.parent.BootstrapProtocol.EC1_MARKER
+                    + b("\n")
+                    + mitogen.parent.BootstrapProtocol.EC2_MARKER
+                    + b("\n"),
+                )
+                self.assertEqual(
+                    b(""),
+                    proc.stderr.read(),
+                )
+            finally:
+                proc.stdout.close()
+                proc.stderr.close()
+
+    def test_stdin_blocking(self):
+        """Test that first stage works with blocking STDIN
+
+        The boot command should read the preamble from STDIN, write all ECO
+        markers to STDOUT, and then execute the preamble.
+
+        This test writes the complete preamble to blocking STDIN.
+
+        1. Fork child reads from blocking STDIN
+        2. Fork child writes all ECO markers to stdout as expected.
+
+        """
+        PIPE_SIZE = 4096
+        # Make sure that not all data can be written in one write operation.
+        CHUNK_SIZE = 2 * PIPE_SIZE
+        r, w = mitogen.core.pipe(blocking=True)
+        with w:
+            try:
+                fcntl.fcntl(w.fileno(), fcntl.F_SETPIPE_SZ, PIPE_SIZE)
+            except AttributeError:
+                pass
+            else:
+                self.assertEqual(fcntl.fcntl(w.fileno(), fcntl.F_GETPIPE_SZ), PIPE_SIZE)
+            proc = subprocess.Popen(
+                args=self.args,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                stdin=r,
+                close_fds=True,
+            )
+            # Close the read fd in the parent
+            r.close()
+
+            view = mitogen.core.BufferType(self.preamble, 0)
+            offset = 0
+            while offset < len(view):
+                end = min(offset + CHUNK_SIZE, len(view))
+                written = os.write(w.fileno(), view[offset:end])
+                self.assertGreater(written, 0)
+                offset += written
+
+        try:
+            try:
+                returncode = proc.wait(timeout=10)
+            except TypeError:
+                returncode = proc.wait()
+        except:
+            proc.kill()
+            self.fail("First stage did not finish")
+        else:
+            try:
+                # proc was killed by SIGTERM?!? TODO Where does this come from?
+                self.assertEqual(-15, proc.returncode)
+                self.assertEqual(
+                    proc.stdout.read(),
+                    mitogen.parent.BootstrapProtocol.EC0_MARKER
+                    + b("\n")
+                    + mitogen.parent.BootstrapProtocol.EC1_MARKER
+                    + b("\n")
+                    + mitogen.parent.BootstrapProtocol.EC2_MARKER
+                    + b("\n"),
+                )
+                self.assertEqual(
+                    b(""),
+                    proc.stderr.read(),
+                )
+            finally:
+                proc.stdout.close()
+                proc.stderr.close()
+
+    def test_eof_too_early(self):
+        """The boot command should write an ECO marker to stdout, read the
+        preamble from stdin, then execute it.
+
+        This test writes some data to STDIN and closes it then to create an
+        EOF situation.
+        1. Fork child tries to read from STDIN, but stops as EOF is received.
+        2. Fork child crashes (trying to decompress the junk data)
+        3. Fork child's file descriptors (write pipes) are closed by the OS
+        4. Fork parent does `dup(<read pipe>, <stdin>)` and `exec(<python>)`
+        5. Python reads `b''` (i.e. EOF) from stdin (a closed pipe)
+        6. Python runs `''` (a valid script) and exits with success"""
 
         proc = subprocess.Popen(
-            args=args,
+            args=self.args,
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             stdin=subprocess.PIPE,
         )
 
         # Do not send all of the data from the preamble
-        proc.stdin.write(preamble[:-128])
+        proc.stdin.write(self.preamble[:-128])
         proc.stdin.flush()
         proc.stdin.close()
         try:
@@ -155,23 +302,22 @@ def test_eof_too_early(self):
                 proc.stderr.close()
 
     def test_timeout_error(self):
-        options = mitogen.parent.Options(max_message_size=123)
-        conn = mitogen.parent.Connection(options, self.router)
-        conn.context = mitogen.core.Context(None, 123)
+        """
+        The boot command should write an ECO marker to stdout, read the
+        preamble from stdin, then execute it.
+
+        This test attaches closes stdin to create a specific failure
+        1. Fork child tries to read from STDIN, but fails as it is closed
+        2. Fork child raises TimeoutError
+        3. Fork child's file descriptors (write pipes) are closed by the OS
+        4. Fork parent does `dup(<read pipe>, <stdin>)` and `exec(<python>)`
+        5. Python reads `b''` (i.e. EOF) from stdin (a closed pipe)
+        6. Python runs `''` (a valid script) and exits with success
+        """
+
         # We do not want to wait the default of 10s, change it to 0.3s
-        conn._first_stage_timeout = 0.3
-        args = conn.get_boot_command()
-
-        # The boot command should write an ECO marker to stdout, read the
-        # preamble from stdin, then execute it.
-
-        # This test attaches closes stdin to create a specific failure
-        # 1. Fork child tries to read from STDIN, but fails as it is closed
-        # 2. Fork child raises TimeoutError
-        # 3. Fork child's file descriptors (write pipes) are closed by the OS
-        # 4. Fork parent does `dup(<read pipe>, <stdin>)` and `exec(<python>)`
-        # 5. Python reads `b''` (i.e. EOF) from stdin (a closed pipe)
-        # 6. Python runs `''` (a valid script) and exits with success
+        self.conn._first_stage_timeout = 0.3
+        args = self.conn.get_boot_command()
 
         proc = subprocess.Popen(
             args=args,