psutil.Popen: inherit from subprocess + support wait(timeout=...) parameter (#1736)

4 files changed, 37 insertions, 16 deletions

diff --git a/psutil/tests/__init__.py b/psutil/tests/__init__.py
index bf3d973f..b31b845d 100644
--- a/psutil/tests/__init__.py
+++ b/psutil/tests/__init__.py
@@ -459,6 +459,20 @@ def terminate(proc_or_pid, sig=signal.SIGTERM, wait_w_timeout=GLOBAL_TIMEOUT):
     """Terminate and flush a psutil.Process, psutil.Popen or
     subprocess.Popen instance.
     """
+    def wait(proc, timeout=None):
+        if sys.version_info < (3, 3) and not \
+                isinstance(proc, (psutil.Process, psutil.Popen)):
+            # subprocess.Popen instance + no timeout arg.
+            try:
+                ret = psutil.Process(proc.pid).wait(timeout)
+            except psutil.NoSuchProcess:
+                pass
+            else:
+                proc.returncode = ret
+                return ret
+        else:
+            return proc.wait(timeout)
+
     if isinstance(proc_or_pid, int):
         try:
             proc = psutil.Process(proc_or_pid)
@@ -467,7 +481,18 @@ def terminate(proc_or_pid, sig=signal.SIGTERM, wait_w_timeout=GLOBAL_TIMEOUT):
     else:
         proc = proc_or_pid
 
-    if isinstance(proc, subprocess.Popen):
+    if isinstance(proc, (psutil.Process, psutil.Popen)):
+        try:
+            proc.send_signal(sig)
+        except psutil.NoSuchProcess:
+            _assert_no_pid(proc.pid)
+        else:
+            if wait_w_timeout:
+                ret = wait(proc, wait_w_timeout)
+                _assert_no_pid(proc.pid)
+                return ret
+    else:
+        # subprocess.Popen instance
         try:
             proc.send_signal(sig)
         except OSError as err:
@@ -475,6 +500,8 @@ def terminate(proc_or_pid, sig=signal.SIGTERM, wait_w_timeout=GLOBAL_TIMEOUT):
                 pass
             elif err.errno != errno.ESRCH:
                 raise
+        except psutil.NoSuchProcess:  # psutil.Popen
+            pass
         if proc.stdout:
             proc.stdout.close()
         if proc.stderr:
@@ -486,17 +513,9 @@ def terminate(proc_or_pid, sig=signal.SIGTERM, wait_w_timeout=GLOBAL_TIMEOUT):
         finally:
             if wait_w_timeout:
                 try:
-                    proc.wait(wait_w_timeout)
+                    return wait(proc, wait_w_timeout)
                 except ChildProcessError:
                     pass
-    else:
-        try:
-            proc.send_signal(sig)
-        except psutil.NoSuchProcess:
-            _assert_no_pid(proc.pid)
-        else:
-            if wait_w_timeout:
-                proc.wait(wait_w_timeout)
                 _assert_no_pid(proc.pid)
 
 
diff --git a/psutil/tests/runner.py b/psutil/tests/runner.py
index bfadc069..d90feabd 100755
--- a/psutil/tests/runner.py
+++ b/psutil/tests/runner.py
@@ -155,7 +155,8 @@ class SuiteLoader:
 
     def from_name(self, name):
         suite = unittest.TestSuite()
-        name = os.path.splitext(os.path.basename(name))[0]
+        if name.endswith('.py'):
+            name = os.path.splitext(os.path.basename(name))[0]
         suite.addTest(unittest.defaultTestLoader.loadTestsFromName(name))
         return suite
 
diff --git a/psutil/tests/test_linux.py b/psutil/tests/test_linux.py
index ea957c17..6d4a934a 100755
--- a/psutil/tests/test_linux.py
+++ b/psutil/tests/test_linux.py
@@ -33,16 +33,16 @@ from psutil.tests import HAS_BATTERY
 from psutil.tests import HAS_CPU_FREQ
 from psutil.tests import HAS_GETLOADAVG
 from psutil.tests import HAS_RLIMIT
-from psutil.tests import SYSMEM_TOLERANCE
 from psutil.tests import mock
 from psutil.tests import PYPY
 from psutil.tests import pyrun
-from psutil.tests import reap_children
 from psutil.tests import reload_module
 from psutil.tests import retry_on_failure
 from psutil.tests import safe_rmpath
 from psutil.tests import sh
 from psutil.tests import skip_on_not_implemented
+from psutil.tests import SYSMEM_TOLERANCE
+from psutil.tests import terminate
 from psutil.tests import ThreadTask
 from psutil.tests import TRAVIS
 from psutil.tests import unittest
@@ -1652,7 +1652,7 @@ class TestProcess(unittest.TestCase):
                 time.sleep(10)
             """ % testfn)
         sproc = pyrun(src)
-        self.addCleanup(reap_children)
+        self.addCleanup(terminate, sproc)
         call_until(lambda: os.listdir('.'), "'%s' not in ret" % testfn)
         p = psutil.Process(sproc.pid)
         time.sleep(.1)
diff --git a/psutil/tests/test_process.py b/psutil/tests/test_process.py
index 8095fc04..ac841eee 100755
--- a/psutil/tests/test_process.py
+++ b/psutil/tests/test_process.py
@@ -1338,7 +1338,7 @@ class TestProcess(unittest.TestCase):
                 pass
 
         zpid = create_zombie_proc()
-        self.addCleanup(reap_children, recursive=True)
+        self.addCleanup(reap_children)
         # A zombie process should always be instantiable
         zproc = psutil.Process(zpid)
         # ...and at least its status always be querable
@@ -1348,7 +1348,7 @@ class TestProcess(unittest.TestCase):
         # ...and as_dict() shouldn't crash
         zproc.as_dict()
         # if cmdline succeeds it should be an empty list
-        ret = succeed_or_zombie_p_exc(zproc.suspend)
+        ret = succeed_or_zombie_p_exc(zproc.cmdline)
         if ret is not None:
             self.assertEqual(ret, [])
 
@@ -1592,6 +1592,7 @@ class TestPopen(unittest.TestCase):
             self.assertTrue(dir(proc))
             self.assertRaises(AttributeError, getattr, proc, 'foo')
             proc.terminate()
+        proc.wait(timeout=3)
 
     def test_ctx_manager(self):
         with psutil.Popen([PYTHON_EXE, "-V"],