Replace psutil usage with subprocess in utils.run_command. (#776)

4 files changed, 35 insertions, 28 deletions

diff --git a/mobly/controllers/android_device_lib/adb.py b/mobly/controllers/android_device_lib/adb.py
index 21019ae..eaaeefb 100644
--- a/mobly/controllers/android_device_lib/adb.py
+++ b/mobly/controllers/android_device_lib/adb.py
@@ -19,7 +19,6 @@ import threading
 import time
 
 from mobly import utils
-import psutil
 
 # Command to use for running ADB commands.
 ADB = 'adb'
@@ -185,7 +184,7 @@ class AdbProxy:
       raise ValueError('Timeout is not a positive value: %s' % timeout)
     try:
       (ret, out, err) = utils.run_command(args, shell=shell, timeout=timeout)
-    except psutil.TimeoutExpired:
+    except subprocess.TimeoutExpired:
       raise AdbTimeoutError(cmd=args, timeout=timeout, serial=self.serial)
 
     if stderr:
diff --git a/mobly/utils.py b/mobly/utils.py
index 53c4750..e3c69aa 100644
--- a/mobly/utils.py
+++ b/mobly/utils.py
@@ -346,23 +346,19 @@ def run_command(cmd,
       std error.
 
   Raises:
-    psutil.TimeoutExpired: The command timed out.
+    subprocess.TimeoutExpired: The command timed out.
   """
-  # Only import psutil when actually needed.
-  # psutil may cause import error in certain env. This way the utils module
-  # doesn't crash upon import.
-  import psutil
   if stdout is None:
     stdout = subprocess.PIPE
   if stderr is None:
     stderr = subprocess.PIPE
-  process = psutil.Popen(cmd,
-                         stdout=stdout,
-                         stderr=stderr,
-                         shell=shell,
-                         cwd=cwd,
-                         env=env,
-                         universal_newlines=universal_newlines)
+  process = subprocess.Popen(cmd,
+                             stdout=stdout,
+                             stderr=stderr,
+                             shell=shell,
+                             cwd=cwd,
+                             env=env,
+                             universal_newlines=universal_newlines)
   timer = None
   timer_triggered = threading.Event()
   if timeout and timeout > 0:
@@ -377,12 +373,15 @@ def run_command(cmd,
     timer.start()
   # If the command takes longer than the timeout, then the timer thread
   # will kill the subprocess, which will make it terminate.
-  (out, err) = process.communicate()
+  out, err = process.communicate()
   if timer is not None:
     timer.cancel()
   if timer_triggered.is_set():
-    raise psutil.TimeoutExpired(timeout, pid=process.pid)
-  return (process.returncode, out, err)
+    raise subprocess.TimeoutExpired(cmd=cwd,
+                                    timeout=timeout,
+                                    output=out,
+                                    stderr=err)
+  return process.returncode, out, err
 
 
 def start_standing_subprocess(cmd, shell=False, env=None):
diff --git a/tests/mobly/controllers/android_device_lib/adb_test.py b/tests/mobly/controllers/android_device_lib/adb_test.py
index fbde202..c5d97d7 100755
--- a/tests/mobly/controllers/android_device_lib/adb_test.py
+++ b/tests/mobly/controllers/android_device_lib/adb_test.py
@@ -147,21 +147,26 @@ class AdbTest(unittest.TestCase):
 
   @mock.patch('mobly.utils.run_command')
   def test_exec_cmd_timed_out(self, mock_run_command):
-    mock_run_command.side_effect = adb.psutil.TimeoutExpired('Timed out')
+    mock_run_command.side_effect = subprocess.TimeoutExpired(cmd='mock_command',
+                                                             timeout=0.01)
     mock_serial = '1234Abcd'
+
     with self.assertRaisesRegex(
         adb.AdbTimeoutError, 'Timed out executing command "adb -s '
         '1234Abcd fake-cmd" after 0.01s.') as context:
       adb.AdbProxy(mock_serial).fake_cmd(timeout=0.01)
+
     self.assertEqual(context.exception.serial, mock_serial)
     self.assertIn(mock_serial, context.exception.cmd)
 
   @mock.patch('mobly.utils.run_command')
   def test_exec_cmd_timed_out_without_serial(self, mock_run_command):
-    mock_run_command.side_effect = adb.psutil.TimeoutExpired('Timed out')
+    mock_run_command.side_effect = subprocess.TimeoutExpired(cmd='mock_command',
+                                                             timeout=0.01)
+
     with self.assertRaisesRegex(
-        adb.AdbTimeoutError, 'Timed out executing command "adb '
-        'fake-cmd" after 0.01s.') as context:
+        adb.AdbTimeoutError,
+        'Timed out executing command "adb fake-cmd" after 0.01s.'):
       adb.AdbProxy().fake_cmd(timeout=0.01)
 
   def test_exec_cmd_with_negative_timeout_value(self):
diff --git a/tests/mobly/utils_test.py b/tests/mobly/utils_test.py
index 42b9418..345336e 100755
--- a/tests/mobly/utils_test.py
+++ b/tests/mobly/utils_test.py
@@ -68,17 +68,19 @@ class UtilsTest(unittest.TestCase):
     self.assertEqual(ret, 0)
 
   def test_run_command_with_timeout_expired(self):
-    with self.assertRaises(psutil.TimeoutExpired):
+    with self.assertRaises(subprocess.TimeoutExpired):
       _ = utils.run_command(self.sleep_cmd(4), timeout=0.01)
 
   @mock.patch('threading.Timer')
-  @mock.patch('psutil.Popen')
+  @mock.patch('subprocess.Popen')
   def test_run_command_with_default_params(self, mock_Popen, mock_Timer):
     mock_command = mock.MagicMock(spec=dict)
     mock_proc = mock_Popen.return_value
     mock_proc.communicate.return_value = ('fake_out', 'fake_err')
     mock_proc.returncode = 0
+
     out = utils.run_command(mock_command)
+
     self.assertEqual(out, (0, 'fake_out', 'fake_err'))
     mock_Popen.assert_called_with(
         mock_command,
@@ -92,7 +94,7 @@ class UtilsTest(unittest.TestCase):
     mock_Timer.assert_not_called()
 
   @mock.patch('threading.Timer')
-  @mock.patch('psutil.Popen')
+  @mock.patch('subprocess.Popen')
   def test_run_command_with_custom_params(self, mock_Popen, mock_Timer):
     mock_command = mock.MagicMock(spec=dict)
     mock_stdout = mock.MagicMock(spec=int)
@@ -104,6 +106,7 @@ class UtilsTest(unittest.TestCase):
     mock_proc = mock_Popen.return_value
     mock_proc.communicate.return_value = ('fake_out', 'fake_err')
     mock_proc.returncode = 127
+
     out = utils.run_command(mock_command,
                             stdout=mock_stdout,
                             stderr=mock_stderr,
@@ -111,6 +114,7 @@ class UtilsTest(unittest.TestCase):
                             timeout=mock_timeout,
                             env=mock_env,
                             universal_newlines=mock_universal_newlines)
+
     self.assertEqual(out, (127, 'fake_out', 'fake_err'))
     mock_Popen.assert_called_with(
         mock_command,
@@ -124,13 +128,13 @@ class UtilsTest(unittest.TestCase):
     mock_Timer.assert_called_with(1234, mock.ANY)
 
   def test_run_command_with_universal_newlines_false(self):
-    (ret, out, err) = utils.run_command(
-        self.sleep_cmd(0.01), universal_newlines=False)
+    (ret, out, err) = utils.run_command(self.sleep_cmd(0.01),
+                                        universal_newlines=False)
     self.assertIsInstance(out, bytes)
 
   def test_run_command_with_universal_newlines_true(self):
-    (ret, out, err) = utils.run_command(
-        self.sleep_cmd(0.01), universal_newlines=True)
+    (ret, out, err) = utils.run_command(self.sleep_cmd(0.01),
+                                        universal_newlines=True)
     self.assertIsInstance(out, str)
 
   def test_start_standing_subproc(self):