Added 'env' and 'ignore_sigint' to RunCommand.

lib/cros_build_lib_unittest.py

 #!/usr/bin/python
 #
 # Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 import errno
 import os
 import shutil
+import signal
 import subprocess
 import tempfile
 import unittest
 import cros_build_lib
 import mox
 
 
 class TestRunCommand(unittest.TestCase):
 
   def setUp(self):
+    # Get the original value for SIGINT so our signal() mock can return the
+    # correct thing.
+    self._old_sigint = signal.getsignal(signal.SIGINT)
+
     self.mox = mox.Mox()
     self.mox.StubOutWithMock(subprocess, 'Popen', use_mock_anything=True)
+    self.mox.StubOutWithMock(signal, 'signal')
     self.proc_mock = self.mox.CreateMockAnything()
     self.error = 'test error'
     self.output = 'test output'
 
   def tearDown(self):
+    # Unset anything that we set with mox.
     self.mox.UnsetStubs()
 
   def _AssertCrEqual(self, expected, actual):
     """Helper method to compare two CommandResult objects.
 
     This is needed since assertEqual does not know how to compare two
     CommandResult objects.
 
     Args:
       expected: a CommandResult object, expected result.
@@ skipping 15 matching lines @@
       rc_kv: key-value pairs passed to RunCommand().
     """
     expected_result = cros_build_lib.CommandResult()
     expected_result.cmd = real_cmd
     expected_result.error = self.error
     expected_result.output = self.output
     if 'exit_code' in rc_kv:
       expected_result.returncode = self.proc_mock.returncode
 
     arg_dict = dict()
-    for attr in 'cwd stdin stdout stderr shell'.split():
+    for attr in 'cwd env stdin stdout stderr shell'.split():
       if attr in sp_kv:
         arg_dict[attr] = sp_kv[attr]
       else:
         if attr == 'shell':
           arg_dict[attr] = False
         else:
           arg_dict[attr] = None
 
+    # If requested, RunCommand will ignore sigints; record that.
+    if rc_kv.get('ignore_sigint'):
+      signal.signal(signal.SIGINT, signal.SIG_IGN).AndReturn(self._old_sigint)
+
     subprocess.Popen(real_cmd, **arg_dict).AndReturn(self.proc_mock)
     self.proc_mock.communicate(None).AndReturn((self.output, self.error))
 
+    # If it ignored them, RunCommand will restore sigints; record that.
+    if rc_kv.get('ignore_sigint'):
+      signal.signal(signal.SIGINT, self._old_sigint).AndReturn(signal.SIG_IGN)
+
     self.mox.ReplayAll()
     actual_result = cros_build_lib.RunCommand(cmd, **rc_kv)
     self.mox.VerifyAll()
 
     self._AssertCrEqual(expected_result, actual_result)
 
-  def testReturnCodeZeroWithArrayCmd(self):
-    """--enter_chroot=False and --cmd is an array of strings."""
+  def testReturnCodeZeroWithArrayCmd(self, ignore_sigint=False):
+    """--enter_chroot=False and --cmd is an array of strings.
+
+    Parameterized so this can also be used by some other tests w/ alternate
+    params to RunCommand().
+
+    Args:
+      ignore_sigint: If True, we'll tell RunCommand to ignore sigint.
+    """
     self.proc_mock.returncode = 0
     cmd_list = ['foo', 'bar', 'roger']
-    self._TestCmd(cmd_list, cmd_list, rc_kv=dict(exit_code=True))
+    self._TestCmd(cmd_list, cmd_list, rc_kv=dict(exit_code=True,
+                                                 ignore_sigint=ignore_sigint))
+
+  def testSignalRestoreNormalCase(self):
+    """Test RunCommand() properly sets/restores sigint.  Normal case."""
+    self.testReturnCodeZeroWithArrayCmd(ignore_sigint=True)
+
 
   def testReturnCodeZeroWithArrayCmdEnterChroot(self):
     """--enter_chroot=True and --cmd is an array of strings."""
     self.proc_mock.returncode = 0
     cmd_list = ['foo', 'bar', 'roger']
     real_cmd = ['./enter_chroot.sh', '--'] + cmd_list
     self._TestCmd(cmd_list, real_cmd, rc_kv=dict(enter_chroot=True))
 
   def testReturnCodeNotZeroErrorOkNotRaisesError(self):
     """Raise error when proc.communicate() returns non-zero."""
     self.proc_mock.returncode = 1
     cmd = 'test cmd'
     self._TestCmd(cmd, cmd, rc_kv=dict(error_ok=True))
 
-  def testSubprocessCommunicateExceptionRaisesError(self):
-    """Verify error raised by communicate() is caught."""
+  def testSubprocessCommunicateExceptionRaisesError(self, ignore_sigint=False):
+    """Verify error raised by communicate() is caught.
+
+    Parameterized so this can also be used by some other tests w/ alternate
+    params to RunCommand().
+
+    Args:
+      ignore_sigint: If True, we'll tell RunCommand to ignore sigint.
+    """
     cmd = 'test cmd'
-    subprocess.Popen(cmd, cwd=None, stdin=None, stdout=None, stderr=None,
+
+    # If requested, RunCommand will ignore sigints; record that.
+    if ignore_sigint:
+      signal.signal(signal.SIGINT, signal.SIG_IGN).AndReturn(self._old_sigint)
+
+    subprocess.Popen(cmd, cwd=None, env=None,
+                     stdin=None, stdout=None, stderr=None,
                      shell=False).AndReturn(self.proc_mock)
     self.proc_mock.communicate(None).AndRaise(ValueError)
 
+    # If it ignored them, RunCommand will restore sigints; record that.
+    if ignore_sigint:
+      signal.signal(signal.SIGINT, self._old_sigint).AndReturn(signal.SIG_IGN)
+
     self.mox.ReplayAll()
-    self.assertRaises(ValueError, cros_build_lib.RunCommand, cmd)
+    self.assertRaises(ValueError, cros_build_lib.RunCommand, cmd,
+                      ignore_sigint=ignore_sigint)
     self.mox.VerifyAll()
 
+  def testSignalRestoreExceptionCase(self):
+    """Test RunCommand() properly sets/restores sigint.  Exception case."""
+    self.testSubprocessCommunicateExceptionRaisesError(ignore_sigint=True)
+
   def testSubprocessCommunicateExceptionNotRaisesError(self):
     """Don't re-raise error from communicate() when --error_ok=True."""
     cmd = 'test cmd'
     real_cmd = './enter_chroot.sh -- %s' % cmd
     expected_result = cros_build_lib.CommandResult()
     expected_result.cmd = real_cmd
 
-    subprocess.Popen(real_cmd, cwd=None, stdin=None, stdout=None, stderr=None,
+    subprocess.Popen(real_cmd, cwd=None, env=None,
+                     stdin=None, stdout=None, stderr=None,
                      shell=False).AndReturn(self.proc_mock)
     self.proc_mock.communicate(None).AndRaise(ValueError)
 
     self.mox.ReplayAll()
     actual_result = cros_build_lib.RunCommand(cmd, error_ok=True,
                                               enter_chroot=True)
     self.mox.VerifyAll()
 
     self._AssertCrEqual(expected_result, actual_result)
 
+  def testEnvWorks(self):
+    """Test RunCommand(..., env=xyz) works."""
+    # We'll put this bogus environment together, just to make sure
+    # subprocess.Popen gets passed it.
+    env = {'Tom': 'Jerry', 'Itchy': 'Scratchy'}
+
+    # This is a simple case, copied from testReturnCodeZeroWithArrayCmd()
+    self.proc_mock.returncode = 0
+    cmd_list = ['foo', 'bar', 'roger']
+
+    # Run.  We expect the env= to be passed through from sp (subprocess.Popen)
+    # to rc (RunCommand).
+    self._TestCmd(cmd_list, cmd_list,
+                  sp_kv=dict(env=env),
+                  rc_kv=dict(env=env, exit_code=True))
+
 
 class TestListFiles(unittest.TestCase):
 
   def setUp(self):
     self.root_dir = tempfile.mkdtemp(prefix='listfiles_unittest')
 
   def tearDown(self):
     shutil.rmtree(self.root_dir)
 
   def _CreateNestedDir(self, dir_structure):
@@ skipping 93 matching lines @@
 
   def testGetChromeosVersionWithEmptyInputReturnsDefault(self):
     self._TestChromeosVersion('')
 
   def testGetChromeosVersionWithNoneInputReturnsDefault(self):
     self._TestChromeosVersion(None)
 
 
 if __name__ == '__main__':
   unittest.main()