From cb7ed2f48c1b7b6e24e896707ea6f4d6b1a2c97d Mon Sep 17 00:00:00 2001
From: Mark Goddard <mark@stackhpc.com>
Date: Thu, 20 Apr 2017 14:29:15 +0100
Subject: [PATCH] Allow utils.run_command to return command output

---
 kayobe/tests/unit/test_utils.py | 27 +++++++++++++++++++--------
 kayobe/utils.py                 | 22 ++++++++++++++++------
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/kayobe/tests/unit/test_utils.py b/kayobe/tests/unit/test_utils.py
index b25cfdba9..2feb93cc1 100644
--- a/kayobe/tests/unit/test_utils.py
+++ b/kayobe/tests/unit/test_utils.py
@@ -70,18 +70,29 @@ key2: value2
 
     @mock.patch.object(subprocess, "check_call")
     def test_run_command(self, mock_call):
-        utils.run_command(["command", "to", "run"])
+        output = utils.run_command(["command", "to", "run"])
         mock_call.assert_called_once_with(["command", "to", "run"])
+        self.assertIsNone(output)
 
+    @mock.patch("kayobe.utils.open")
     @mock.patch.object(subprocess, "check_call")
-    def test_run_command_quiet(self, mock_call):
-        utils.run_command(["command", "to", "run"], quiet=True)
+    def test_run_command_quiet(self, mock_call, mock_open):
+        mock_devnull = mock_open.return_value.__enter__.return_value
+        output = utils.run_command(["command", "to", "run"], quiet=True)
         mock_call.assert_called_once_with(["command", "to", "run"],
-                                          stdout=subprocess.PIPE,
-                                          stderr=subprocess.PIPE)
+                                          stdout=mock_devnull,
+                                          stderr=mock_devnull)
+        self.assertIsNone(output)
 
-    @mock.patch.object(subprocess, "check_call")
-    def test_run_command_failure(self, mock_call):
-        mock_call.side_effect = subprocess.CalledProcessError(1, "command")
+    @mock.patch.object(subprocess, "check_output")
+    def test_run_command_check_output(self, mock_output):
+        mock_output.return_value = "command output"
+        output = utils.run_command(["command", "to", "run"], check_output=True)
+        mock_output.assert_called_once_with(["command", "to", "run"])
+        self.assertEqual(output, "command output")
+
+    @mock.patch.object(subprocess, "check_output")
+    def test_run_command_failure(self, mock_output):
+        mock_output.side_effect = subprocess.CalledProcessError(1, "command")
         self.assertRaises(subprocess.CalledProcessError, utils.run_command,
                           ["command", "to", "run"])
diff --git a/kayobe/utils.py b/kayobe/utils.py
index 91376bb36..08d4c1d0b 100644
--- a/kayobe/utils.py
+++ b/kayobe/utils.py
@@ -93,14 +93,24 @@ def is_readable_file(path):
     return {"result": True}
 
 
-def run_command(cmd, quiet=False, **kwargs):
-    """Run a command, checking the output."""
-    if quiet:
-        kwargs["stdout"] = subprocess.PIPE
-        kwargs["stderr"] = subprocess.PIPE
+def run_command(cmd, quiet=False, check_output=False, **kwargs):
+    """Run a command, checking the output.
+
+    :param quiet: Redirect output to /dev/null
+    :param check_output: Whether to return the output of the command
+    :returns: The output of the command if check_output is true
+    """
     if isinstance(cmd, six.string_types):
         cmd_string = cmd
     else:
         cmd_string = " ".join(cmd)
     LOG.debug("Running command: %s", cmd_string)
-    subprocess.check_call(cmd, **kwargs)
+    if quiet:
+        with open("/dev/null", "w") as devnull:
+            kwargs["stdout"] = devnull
+            kwargs["stderr"] = devnull
+            subprocess.check_call(cmd, **kwargs)
+    elif check_output:
+        return subprocess.check_output(cmd, **kwargs)
+    else:
+        subprocess.check_call(cmd, **kwargs)