Refactor set_forkserver_preload to use _handle_preload helper

Extract preload logic into a separate _handle_preload() function to
enable targeted unit testing. Add comprehensive unit tests for both
module and __main__ preload with all three on_error modes.

New tests:
- test_handle_preload_main_on_error_{fail,warn,ignore}
- test_handle_preload_module_on_error_{fail,warn,ignore}
- test_handle_preload_main_valid
- test_handle_preload_combined

Total test count increased from 6 to 14 tests.

Author: gpshead

Lib/multiprocessing/forkserver.py

@@ -219,6 +219,67 @@ def ensure_running(self):
 #
 #
 
+def _handle_preload(preload, main_path=None, sys_path=None, on_error='ignore'):
+    """Handle module preloading with configurable error handling.
+
+    Args:
+        preload: List of module names to preload.
+        main_path: Path to __main__ module if '__main__' is in preload.
+        sys_path: sys.path to use for imports (None means use current).
+        on_error: How to handle import errors ('ignore', 'warn', or 'fail').
+    """
+    if not preload:
+        return
+
+    if sys_path is not None:
+        sys.path[:] = sys_path
+
+    if '__main__' in preload and main_path is not None:
+        process.current_process()._inheriting = True
+        try:
+            spawn.import_main_path(main_path)
+        except Exception as e:
+            match on_error:
+                case 'fail':
+                    raise
+                case 'warn':
+                    import warnings
+                    with warnings.catch_warnings():
+                        warnings.simplefilter('always', ImportWarning)
+                        warnings.warn(
+                            f"Failed to preload __main__ from {main_path!r}: {e}",
+                            ImportWarning,
+                            stacklevel=2
+                        )
+                case 'ignore':
+                    pass
+        finally:
+            del process.current_process()._inheriting
+
+    for modname in preload:
+        try:
+            __import__(modname)
+        except ImportError as e:
+            match on_error:
+                case 'fail':
+                    raise
+                case 'warn':
+                    import warnings
+                    with warnings.catch_warnings():
+                        warnings.simplefilter('always', ImportWarning)
+                        warnings.warn(
+                            f"Failed to preload module {modname!r}: {e}",
+                            ImportWarning,
+                            stacklevel=2
+                        )
+                case 'ignore':
+                    pass
+
+    # gh-135335: flush stdout/stderr in case any of the preloaded modules
+    # wrote to them, otherwise children might inherit buffered data
+    util._flush_std_streams()
+
+
 def main(listener_fd, alive_r, preload, main_path=None, sys_path=None,
          *, authkey_r=None, on_error='ignore'):
     """Run forkserver."""
@@ -231,52 +292,7 @@ def main(listener_fd, alive_r, preload, main_path=None, sys_path=None,
     else:
         authkey = b''
 
-    if preload:
-        if sys_path is not None:
-            sys.path[:] = sys_path
-        if '__main__' in preload and main_path is not None:
-            process.current_process()._inheriting = True
-            try:
-                spawn.import_main_path(main_path)
-            except Exception as e:
-                match on_error:
-                    case 'fail':
-                        raise
-                    case 'warn':
-                        import warnings
-                        with warnings.catch_warnings():
-                            warnings.simplefilter('always', ImportWarning)
-                            warnings.warn(
-                                f"Failed to preload __main__ from {main_path!r}: {e}",
-                                ImportWarning,
-                                stacklevel=2
-                            )
-                    case 'ignore':
-                        pass
-            finally:
-                del process.current_process()._inheriting
-        for modname in preload:
-            try:
-                __import__(modname)
-            except ImportError as e:
-                match on_error:
-                    case 'fail':
-                        raise
-                    case 'warn':
-                        import warnings
-                        with warnings.catch_warnings():
-                            warnings.simplefilter('always', ImportWarning)
-                            warnings.warn(
-                                f"Failed to preload module {modname!r}: {e}",
-                                ImportWarning,
-                                stacklevel=2
-                            )
-                    case 'ignore':
-                        pass
-
-        # gh-135335: flush stdout/stderr in case any of the preloaded modules
-        # wrote to them, otherwise children might inherit buffered data
-        util._flush_std_streams()
+    _handle_preload(preload, main_path, sys_path, on_error)
 
     util._close_stdin()
 

Lib/test/test_multiprocessing_forkserver/test_preload.py

@@ -2,7 +2,10 @@
 
 import multiprocessing
 import multiprocessing.forkserver
+import os
+import tempfile
 import unittest
+from multiprocessing.forkserver import _handle_preload
 
 
 class TestForkserverPreload(unittest.TestCase):
@@ -116,5 +119,88 @@ def test_preload_invalid_on_error_value(self):
         self.assertIn("on_error must be 'ignore', 'warn', or 'fail'", str(cm.exception))
 
 
+class TestHandlePreload(unittest.TestCase):
+    """Unit tests for _handle_preload() function."""
+
+    def test_handle_preload_main_on_error_fail(self):
+        """Test that __main__ import failures raise with on_error='fail'."""
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
+            f.write('raise RuntimeError("test error in __main__")\n')
+            bad_main_path = f.name
+
+        try:
+            with self.assertRaises(RuntimeError) as cm:
+                _handle_preload(['__main__'], main_path=bad_main_path, on_error='fail')
+            self.assertIn("test error in __main__", str(cm.exception))
+        finally:
+            os.unlink(bad_main_path)
+
+    def test_handle_preload_main_on_error_warn(self):
+        """Test that __main__ import failures warn with on_error='warn'."""
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
+            f.write('raise ImportError("test import error")\n')
+            bad_main_path = f.name
+
+        try:
+            with self.assertWarns(ImportWarning) as cm:
+                _handle_preload(['__main__'], main_path=bad_main_path, on_error='warn')
+            self.assertIn("Failed to preload __main__", str(cm.warning))
+            self.assertIn("test import error", str(cm.warning))
+        finally:
+            os.unlink(bad_main_path)
+
+    def test_handle_preload_main_on_error_ignore(self):
+        """Test that __main__ import failures are ignored with on_error='ignore'."""
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
+            f.write('raise ImportError("test import error")\n')
+            bad_main_path = f.name
+
+        try:
+            # Should not raise
+            _handle_preload(['__main__'], main_path=bad_main_path, on_error='ignore')
+        finally:
+            os.unlink(bad_main_path)
+
+    def test_handle_preload_main_valid(self):
+        """Test that valid __main__ preload works."""
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
+            f.write('test_var = 42\n')
+            valid_main_path = f.name
+
+        try:
+            _handle_preload(['__main__'], main_path=valid_main_path, on_error='fail')
+            # Should complete without raising
+        finally:
+            os.unlink(valid_main_path)
+
+    def test_handle_preload_module_on_error_fail(self):
+        """Test that module import failures raise with on_error='fail'."""
+        with self.assertRaises(ModuleNotFoundError):
+            _handle_preload(['nonexistent_test_module_xyz'], on_error='fail')
+
+    def test_handle_preload_module_on_error_warn(self):
+        """Test that module import failures warn with on_error='warn'."""
+        with self.assertWarns(ImportWarning) as cm:
+            _handle_preload(['nonexistent_test_module_xyz'], on_error='warn')
+        self.assertIn("Failed to preload module", str(cm.warning))
+
+    def test_handle_preload_module_on_error_ignore(self):
+        """Test that module import failures are ignored with on_error='ignore'."""
+        # Should not raise
+        _handle_preload(['nonexistent_test_module_xyz'], on_error='ignore')
+
+    def test_handle_preload_combined(self):
+        """Test preloading both __main__ and modules."""
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
+            f.write('import sys\n')
+            valid_main_path = f.name
+
+        try:
+            _handle_preload(['__main__', 'os', 'sys'], main_path=valid_main_path, on_error='fail')
+            # Should complete without raising
+        finally:
+            os.unlink(valid_main_path)
+
+
 if __name__ == '__main__':
     unittest.main()