command_executer: Allow passing custom os env.

It is currently tedious to pass custom env variables to
command executer e.g. when there is a need to pass
an env variable for use inside chroot command.
Allow passing an env variable directly to command_executer.

Includes formatting changes done by running yapf.

BUG=None
TEST=unit tests

Change-Id: I1bd286b89f3938c1c2e8384c2b9f3bdea05f3dd8
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/1733646
Reviewed-by: George Burgess <gbiv@chromium.org>
Tested-by: Manoj Gupta <manojgupta@chromium.org>

1 files changed, 23 insertions, 18 deletions

diff --git a/cros_utils/command_executer.py b/cros_utils/command_executer.py
index ae1b2962..3e710afe 100644
--- a/cros_utils/command_executer.py
+++ b/cros_utils/command_executer.py
@@ -1,6 +1,7 @@
 # Copyright 2011 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.
+
 """Utilities to run commands in outside/inside chroot and on the board."""
 
 from __future__ import print_function
@@ -66,6 +67,7 @@ class CommandExecuter(object):
                         command_timeout=None,
                         terminated_timeout=10,
                         print_to_console=True,
+                        env=None,
                         except_handler=lambda p, e: None):
     """Run a command.
 
@@ -104,7 +106,8 @@ class CommandExecuter(object):
           stderr=subprocess.PIPE,
           shell=True,
           preexec_fn=os.setsid,
-          executable='/bin/bash')
+          executable='/bin/bash',
+          env=env)
 
       full_stdout = ''
       full_stderr = ''
@@ -124,9 +127,9 @@ class CommandExecuter(object):
         if command_terminator and command_terminator.IsTerminated():
           os.killpg(os.getpgid(p.pid), signal.SIGTERM)
           if self.logger:
-            self.logger.LogError('Command received termination request. '
-                                 'Killed child process group.',
-                                 print_to_console)
+            self.logger.LogError(
+                'Command received termination request. '
+                'Killed child process group.', print_to_console)
           break
 
         l = my_poll.poll(100)
@@ -156,18 +159,19 @@ class CommandExecuter(object):
           elif (terminated_timeout is not None and
                 time.time() - terminated_time > terminated_timeout):
             if self.logger:
-              self.logger.LogWarning('Timeout of %s seconds reached since '
-                                     'process termination.' %
-                                     terminated_timeout, print_to_console)
+              self.logger.LogWarning(
+                  'Timeout of %s seconds reached since '
+                  'process termination.' % terminated_timeout, print_to_console)
             break
 
         if (command_timeout is not None and
             time.time() - started_time > command_timeout):
           os.killpg(os.getpgid(p.pid), signal.SIGTERM)
           if self.logger:
-            self.logger.LogWarning('Timeout of %s seconds reached since process'
-                                   'started. Killed child process group.' %
-                                   command_timeout, print_to_console)
+            self.logger.LogWarning(
+                'Timeout of %s seconds reached since process'
+                'started. Killed child process group.' % command_timeout,
+                print_to_console)
           break
 
         if out == err == '':
@@ -371,8 +375,8 @@ class CommandExecuter(object):
     # the chroot already exists. We want the final returned output to skip
     # the output from chroot creation steps.
     if return_output:
-      ret = self.RunCommand('cd %s; cros_sdk %s -- true' % (chromeos_root,
-                                                            cros_sdk_options))
+      ret = self.RunCommand(
+          'cd %s; cros_sdk %s -- true' % (chromeos_root, cros_sdk_options))
       if ret:
         return (ret, '', '')
 
@@ -449,15 +453,15 @@ class CommandExecuter(object):
       src = src + '/'
       dest = dest + '/'
 
-    if src_cros == True or dest_cros == True:
+    if src_cros or dest_cros:
       if self.logger:
-        self.logger.LogFatalIf(src_cros == dest_cros,
-                               'Only one of src_cros and desc_cros can '
-                               'be True.')
+        self.logger.LogFatalIf(
+            src_cros == dest_cros, 'Only one of src_cros and desc_cros can '
+            'be True.')
         self.logger.LogFatalIf(not chromeos_root, 'chromeos_root not given!')
       elif src_cros == dest_cros or not chromeos_root:
         sys.exit(1)
-      if src_cros == True:
+      if src_cros:
         cros_machine = src_machine
       else:
         cros_machine = dest_machine
@@ -467,7 +471,7 @@ class CommandExecuter(object):
           'ssh -p ${FLAGS_ssh_port}' + ' -o StrictHostKeyChecking=no' +
           ' -o UserKnownHostsFile=$(mktemp)' + ' -i $TMP_PRIVATE_KEY')
       rsync_prefix = "\nrsync -r -e \"%s\" " % ssh_command
-      if dest_cros == True:
+      if dest_cros:
         command += rsync_prefix + '%s root@%s:%s' % (src, dest_machine, dest)
         return self.RunCommand(
             command,
@@ -654,6 +658,7 @@ class MockCommandExecuter(CommandExecuter):
                         command_timeout=None,
                         terminated_timeout=10,
                         print_to_console=True,
+                        env=None,
                         except_handler=lambda p, e: None):
     assert not command_timeout
     cmd = str(cmd)