bpo-30599: Fix test_threaded_import reference leak (#2029)

Mock os.register_at_fork() when importing the random module, since
this function doesn't allow to unregister callbacks and so leaked
memory.

1 files changed, 11 insertions, 2 deletions

diff --git a/Lib/test/test_threaded_import.py b/Lib/test/test_threaded_import.py
index 9b2d9a6..a4cf1d7 100644
--- a/Lib/test/test_threaded_import.py
+++ b/Lib/test/test_threaded_import.py
@@ -12,6 +12,7 @@ import sys
 import time
 import shutil
 import unittest
+from unittest import mock
 from test.support import (
     verbose, import_module, run_unittest, TESTFN, reap_threads,
     forget, unlink, rmtree, start_threads)
@@ -37,6 +38,12 @@ def task(N, done, done_tasks, errors):
         if finished:
             done.set()
 
+def mock_register_at_fork(func):
+    # bpo-30599: Mock os.register_at_fork() when importing the random module,
+    # since this function doesn't allow to unregister callbacks and would leak
+    # memory.
+    return mock.patch('os.register_at_fork', create=True)(func)
+
 # Create a circular import structure: A -> C -> B -> D -> A
 # NOTE: `time` is already loaded and therefore doesn't threaten to deadlock.
 
@@ -97,7 +104,8 @@ class ThreadedImportTests(unittest.TestCase):
         if self.old_random is not None:
             sys.modules['random'] = self.old_random
 
-    def check_parallel_module_init(self):
+    @mock_register_at_fork
+    def check_parallel_module_init(self, mock_os):
         if imp.lock_held():
             # This triggers on, e.g., from test import autotest.
             raise unittest.SkipTest("can't run when import lock is held")
@@ -214,7 +222,8 @@ class ThreadedImportTests(unittest.TestCase):
         t2.join()
         self.assertEqual(set(results), {'a', 'b'})
 
-    def test_side_effect_import(self):
+    @mock_register_at_fork
+    def test_side_effect_import(self, mock_os):
         code = """if 1:
             import threading
             def target():