Merge pull request #5 from rezib/pr/run-yoann-comments

Some improvements in run module

Author: valeriyoann

lib/rift/Mock.py

@@ -172,7 +172,7 @@ def _exec(self, cmd):
             cmd,
             live_output=logging.getLogger().isEnabledFor(logging.INFO),
             capture_output=True,
-            merged_capture=True,
+            merge_out_err=True,
             cwd='/'
         )
         if proc.returncode != 0:

lib/rift/run.py

@@ -45,8 +45,30 @@
     'RunResult', ['returncode', 'out', 'err']
 )
 
-def _handle_process_output(process, selector):
-    """Handle process output events with registered selectors."""
+def _handle_process_output(process, live_output, buf_out, buf_err):
+    """Handle process output until it is terminated."""
+
+    # Process output lines handlers
+    def handle_stdout_line(line):
+        buf_out.write(line)
+        if live_output:
+            sys.stdout.write(line)
+    def handle_stderr_line(line):
+        buf_err.write(line)
+        if live_output:
+            sys.stderr.write(line)
+
+    # Process output event handlers
+    def handle_stdout_event(stream):
+        handle_stdout_line(stream.readline())
+    def handle_stderr_event(stream):
+        handle_stderr_line(stream.readline())
+
+    # Register callback for read events from subprocess stdout/stderr streams
+    selector = selectors.DefaultSelector()
+    selector.register(process.stdout, selectors.EVENT_READ, handle_stdout_event)
+    selector.register(process.stderr, selectors.EVENT_READ, handle_stderr_event)
+
     # Loop until subprocess is terminated
     while process.poll() is None:
         # Wait for events and handle them with their registered callbacks
@@ -55,18 +77,31 @@ def _handle_process_output(process, selector):
             callback = key.data
             callback(key.fileobj)
 
+    # Close selector
+    selector.close()
+
+    # The loop above stops processing output as soon as the process is
+    # terminated. However, there may still be buffered output to flush.
+    for line in process.stdout:
+        handle_stdout_line(line)
+    for line in process.stderr:
+        handle_stderr_line(line)
+
+    # Ensure process is terminated
+    process.wait()
+
 def run_command(
         cmd,
         live_output=True,
         capture_output=False,
-        merged_capture=False,
+        merge_out_err=False,
         **kwargs
     ):
     """
-    Run a command and return RunResult named tuple. When live_output is True,
-    command stdout/stderr are streamed in live in current process stdout/stderr.
-    When capture_output is True, command stdout/stderr are available in out/err
-    attributes of RunResult namedtuple. When merged_capture is True as well,
+    Run a command and return a RunResult named tuple. When live_output is True,
+    command stdout/stderr are redirected to current process stdout/stderr. When
+    capture_output is True, command stdout/stderr are available in out/err
+    attributes of RunResult namedtuple. When merge_out_err is True as well,
     command stderr is merged with stdout in out attribute of RunResult named
     tuple. In this case, err attribute is None.
 
@@ -100,55 +135,18 @@ def run_command(
         # Initialize string buffers to store process output in memory
         buf_out = io.StringIO()
         buf_err = None
-        if not merged_capture:
+        if merge_out_err:
+            buf_err = buf_out
+        else:
             buf_err = io.StringIO()
 
-        # Process output lines handlers
-        def handle_stdout(stream):
-            line = stream.readline()
-            buf_out.write(line)
-            if live_output:
-                sys.stdout.write(line)
-        def handle_stderr(stream):
-            line = stream.readline()
-            if merged_capture:
-                buf_out.write(line)
-            else:
-                buf_err.write(line)
-            if live_output:
-                sys.stderr.write(line)
-
-        # Register callback for read events from subprocess stdout/stderr streams
-        selector = selectors.DefaultSelector()
-        selector.register(process.stdout, selectors.EVENT_READ, handle_stdout)
-        selector.register(process.stderr, selectors.EVENT_READ, handle_stderr)
-
-        # Handle process output with registered selectors
-        _handle_process_output(process, selector)
-
-        # _handle_process_output stops processing output as soon as the process
-        # is terminated. However, there may still be buffered output to flush.
-        for line in process.stdout:
-            buf_out.write(line)
-            if live_output:
-                sys.stdout.write(line)
-        for line in process.stderr:
-            if merged_capture:
-                buf_out.write(line)
-            else:
-                buf_err.write(line)
-            if live_output:
-                sys.stdout.write(line)
-
-        # Ensure process is terminated
-        process.wait()
-
-        # Store buffered output
-        selector.close()
+        # Handle process output
+        _handle_process_output(process, live_output, buf_out, buf_err)
 
+    # Get values for out/err buffers and close them
     out = buf_out.getvalue()
     buf_out.close()
-    if merged_capture:
+    if merge_out_err:
         err = None
     else:
         err = buf_err.getvalue()

tests/run.py

@@ -25,13 +25,16 @@ def test_run_command_failed(self):
         """ Test run_command() with basic failed command. """
         proc = run_command(["/bin/false"])
         self.assertEqual(proc.returncode, 1)
+        self.assertIsNone(proc.out)
+        self.assertIsNone(proc.err)
 
     @patch('sys.stdout', new_callable=StringIO)
     def test_run_command_capture_stdout(self, mock_stdout):
         """ Test run_command() with captured standard output. """
         proc = run_command(["/bin/echo", "output_data"], capture_output=True)
         # Standard output must be available in out attribute of RunResult named
         # tuple.
+        self.assertEqual(proc.returncode, 0)
         self.assertEqual(proc.out, "output_data\n")
         self.assertEqual(proc.err, "")
         # Standard output must also be streamed into current process stdout.
@@ -43,26 +46,42 @@ def test_run_command_capture_stderr(self, mock_stderr):
         proc = run_command("/bin/echo error_data 1>&2", capture_output=True, shell=True)
         # Standard err must be available in err attribute of RunResult named
         # tuple.
+        self.assertEqual(proc.returncode, 0)
         self.assertEqual(proc.out, "")
         self.assertEqual(proc.err, "error_data\n")
         # Standard error must also be streamed into current process stderr.
         self.assertEqual(mock_stderr.getvalue(), "error_data\n")
 
     @patch('sys.stderr', new_callable=StringIO)
-    def test_run_command_capture_merged(self, mock_stderr):
-        """ Test run_command() with merged output capture. """
-        proc = run_command("/bin/echo error_data 1>&2", capture_output=True, merged_capture=True, shell=True)
+    def test_run_command_capture_stderr_merged(self, mock_stderr):
+        """ Test run_command() with merged error output capture. """
+        proc = run_command("/bin/echo error_data 1>&2", capture_output=True,
+                merge_out_err=True, shell=True)
         # With merged_capture, standard err must be available in out attribute
         # of RunResult named tuple, and err attribute must be None.
         self.assertEqual(proc.out, "error_data\n")
         self.assertIsNone(proc.err)
         # Standard error must also be streamed into current process stderr.
         self.assertEqual(mock_stderr.getvalue(), "error_data\n")
 
+    @patch('sys.stderr', new_callable=StringIO)
+    def test_run_command_capture_both_merged(self, mock_stderr):
+        """ Test run_command() with merged error and standard output capture. """
+        proc = run_command("/bin/echo error_data 1>&2 && /bin/echo output_data",
+                capture_output=True, merge_out_err=True, shell=True)
+        # With merge_out_err, standard err must be available in out attribute
+        # of RunResult named tuple, and err attribute must be None.
+        self.assertEqual(proc.out, "error_data\noutput_data\n")
+        self.assertIsNone(proc.err)
+        # Standard error must also be streamed into current process stderr.
+        self.assertEqual(mock_stderr.getvalue(), "error_data\n")
+
+    @patch('sys.stderr', new_callable=StringIO)
     @patch('sys.stdout', new_callable=StringIO)
-    def test_run_command_no_output(self, mock_stdout):
+    def test_run_command_no_output(self, mock_stdout, mock_stderr):
         """ Test run_command() without live output. """
-        # With live_output disabled, standard output must not be streamed into
-        # current process stdout.
+        # With live_output disabled, standard output and standard error must not
+        # be redirected in current process stdout.
         proc = run_command(["/bin/echo", "output_data"], live_output=False)
         self.assertEqual(mock_stdout.getvalue(), "")
+        self.assertEqual(mock_stderr.getvalue(), "")