Closed pool when parallel test runner encounters unpicklable exceptions.

2 files changed, 29 insertions, 31 deletions

diff --git a/django/test/runner.py b/django/test/runner.py
index ecae164d7f..5cd72119f5 100644
--- a/django/test/runner.py
+++ b/django/test/runner.py
@@ -567,7 +567,14 @@ class ParallelTestSuite(unittest.TestSuite):
         """
         self.initialize_suite()
         counter = multiprocessing.Value(ctypes.c_int, 0)
-        pool = multiprocessing.Pool(
+        args = [
+            (self.runner_class, index, subsuite, self.failfast, self.buffer)
+            for index, subsuite in enumerate(self.subsuites)
+        ]
+        # Don't buffer in the main process to avoid error propagation issues.
+        result.buffer = False
+
+        with multiprocessing.Pool(
             processes=self.processes,
             initializer=functools.partial(_safe_init_worker, self.init_worker.__func__),
             initargs=[
@@ -579,38 +586,30 @@ class ParallelTestSuite(unittest.TestSuite):
                 self.debug_mode,
                 self.used_aliases,
             ],
-        )
-        args = [
-            (self.runner_class, index, subsuite, self.failfast, self.buffer)
-            for index, subsuite in enumerate(self.subsuites)
-        ]
-        # Don't buffer in the main process to avoid error propagation issues.
-        result.buffer = False
+        ) as pool:
+            test_results = pool.imap_unordered(self.run_subsuite.__func__, args)
 
-        test_results = pool.imap_unordered(self.run_subsuite.__func__, args)
+            while True:
+                if result.shouldStop:
+                    pool.terminate()
+                    break
 
-        while True:
-            if result.shouldStop:
-                pool.terminate()
-                break
-
-            try:
-                subsuite_index, events = test_results.next(timeout=0.1)
-            except multiprocessing.TimeoutError as err:
-                if counter.value < 0:
-                    err.add_note("ERROR: _init_worker failed, see prior traceback")
+                try:
+                    subsuite_index, events = test_results.next(timeout=0.1)
+                except multiprocessing.TimeoutError as err:
+                    if counter.value < 0:
+                        err.add_note("ERROR: _init_worker failed, see prior traceback")
+                        raise
+                    continue
+                except StopIteration:
                     pool.close()
-                    raise
-                continue
-            except StopIteration:
-                pool.close()
-                break
+                    break
 
-            tests = list(self.subsuites[subsuite_index])
-            for event in events:
-                self.handle_event(result, tests, event)
+                tests = list(self.subsuites[subsuite_index])
+                for event in events:
+                    self.handle_event(result, tests, event)
 
-        pool.join()
+            pool.join()
 
         return result
 
diff --git a/tests/test_runner/test_parallel.py b/tests/test_runner/test_parallel.py
index 32cc971d30..193afea1cc 100644
--- a/tests/test_runner/test_parallel.py
+++ b/tests/test_runner/test_parallel.py
@@ -309,9 +309,8 @@ class ParallelTestSuiteTest(SimpleTestCase):
                 test_result.shouldStop = True
                 return (0, remote_result.events)
 
-            mock_pool.return_value.imap_unordered.return_value = unittest.mock.Mock(
-                next=fake_next
-            )
+            mock_imap = mock_pool.return_value.__enter__.return_value.imap_unordered
+            mock_imap.return_value = unittest.mock.Mock(next=fake_next)
             pts.run(test_result)
 
         self.assertIn("ValueError: woops", test_result.errors[0][1])