bpo-40094: Enhance fork and wait tests (GH-19259)

* test_fork1: remove duplicated wait_impl() method: reuse
  fork_wait.py implementation instead.
* Use exit code different than 0 to ensure that we executed the
  expected code path.

4 files changed, 15 insertions, 24 deletions

diff --git a/Lib/test/fork_wait.py b/Lib/test/fork_wait.py
index 8c177556a4..249b5e9607 100644
--- a/Lib/test/fork_wait.py
+++ b/Lib/test/fork_wait.py
@@ -43,8 +43,8 @@ class ForkWait(unittest.TestCase):
             except OSError:
                 pass
 
-    def wait_impl(self, cpid):
-        support.wait_process(cpid, exitcode=0)
+    def wait_impl(self, cpid, *, exitcode):
+        support.wait_process(cpid, exitcode=exitcode)
 
     def test_wait(self):
         for i in range(NUM_THREADS):
@@ -79,4 +79,4 @@ class ForkWait(unittest.TestCase):
             os._exit(n)
         else:
             # Parent
-            self.wait_impl(cpid)
+            self.wait_impl(cpid, exitcode=0)
diff --git a/Lib/test/test_fork1.py b/Lib/test/test_fork1.py
index ce0a05a541..a2f7cfee9c 100644
--- a/Lib/test/test_fork1.py
+++ b/Lib/test/test_fork1.py
@@ -17,19 +17,6 @@ from test import support
 support.get_attribute(os, 'fork')
 
 class ForkTest(ForkWait):
-    def wait_impl(self, cpid):
-        deadline = time.monotonic() + support.SHORT_TIMEOUT
-        while time.monotonic() <= deadline:
-            # waitpid() shouldn't hang, but some of the buildbots seem to hang
-            # in the forking tests.  This is an attempt to fix the problem.
-            spid, status = os.waitpid(cpid, os.WNOHANG)
-            if spid == cpid:
-                break
-            time.sleep(0.1)
-
-        self.assertEqual(spid, cpid)
-        self.assertEqual(status, 0, "cause = %d, exit = %d" % (status&0xff, status>>8))
-
     def test_threaded_import_lock_fork(self):
         """Check fork() in main thread works while a subthread is doing an import"""
         import_started = threading.Event()
@@ -46,6 +33,7 @@ class ForkTest(ForkWait):
         t = threading.Thread(target=importer)
         t.start()
         import_started.wait()
+        exitcode = 42
         pid = os.fork()
         try:
             # PyOS_BeforeFork should have waited for the import to complete
@@ -54,7 +42,7 @@ class ForkTest(ForkWait):
             if not pid:
                 m = __import__(fake_module_name)
                 if m == complete_module:
-                    os._exit(0)
+                    os._exit(exitcode)
                 else:
                     if support.verbose > 1:
                         print("Child encountered partial module")
@@ -64,7 +52,7 @@ class ForkTest(ForkWait):
                 # Exitcode 1 means the child got a partial module (bad.) No
                 # exitcode (but a hang, which manifests as 'got pid 0')
                 # means the child deadlocked (also bad.)
-                self.wait_impl(pid)
+                self.wait_impl(pid, exitcode=exitcode)
         finally:
             try:
                 os.kill(pid, signal.SIGKILL)
@@ -74,6 +62,7 @@ class ForkTest(ForkWait):
 
     def test_nested_import_lock_fork(self):
         """Check fork() in main thread works while the main thread is doing an import"""
+        exitcode = 42
         # Issue 9573: this used to trigger RuntimeError in the child process
         def fork_with_import_lock(level):
             release = 0
@@ -95,8 +84,8 @@ class ForkTest(ForkWait):
                     os._exit(1)
                 raise
             if in_child:
-                os._exit(0)
-            self.wait_impl(pid)
+                os._exit(exitcode)
+            self.wait_impl(pid, exitcode=exitcode)
 
         # Check this works with various levels of nested
         # import in the main thread
diff --git a/Lib/test/test_wait3.py b/Lib/test/test_wait3.py
index 2dc63aaffa..6e06049fbd 100644
--- a/Lib/test/test_wait3.py
+++ b/Lib/test/test_wait3.py
@@ -16,7 +16,7 @@ if not hasattr(os, 'wait3'):
     raise unittest.SkipTest("os.wait3 not defined")
 
 class Wait3Test(ForkWait):
-    def wait_impl(self, cpid):
+    def wait_impl(self, cpid, *, exitcode):
         # This many iterations can be required, since some previously run
         # tests (e.g. test_ctypes) could have spawned a lot of children
         # very quickly.
@@ -30,7 +30,8 @@ class Wait3Test(ForkWait):
             time.sleep(0.1)
 
         self.assertEqual(spid, cpid)
-        self.assertEqual(status, 0, "cause = %d, exit = %d" % (status&0xff, status>>8))
+        self.assertEqual(status, exitcode << 8,
+                         "cause = %d, exit = %d" % (status&0xff, status>>8))
         self.assertTrue(rusage)
 
     def test_wait3_rusage_initialized(self):
diff --git a/Lib/test/test_wait4.py b/Lib/test/test_wait4.py
index a18607252a..6c7ebcb3dd 100644
--- a/Lib/test/test_wait4.py
+++ b/Lib/test/test_wait4.py
@@ -14,7 +14,7 @@ support.get_attribute(os, 'wait4')
 
 
 class Wait4Test(ForkWait):
-    def wait_impl(self, cpid):
+    def wait_impl(self, cpid, *, exitcode):
         option = os.WNOHANG
         if sys.platform.startswith('aix'):
             # Issue #11185: wait4 is broken on AIX and will always return 0
@@ -29,7 +29,8 @@ class Wait4Test(ForkWait):
                 break
             time.sleep(0.1)
         self.assertEqual(spid, cpid)
-        self.assertEqual(status, 0, "cause = %d, exit = %d" % (status&0xff, status>>8))
+        self.assertEqual(status, exitcode << 8,
+                         "cause = %d, exit = %d" % (status&0xff, status>>8))
         self.assertTrue(rusage)
 
 def tearDownModule():