From c28a6fcb3f0b5e4f2fbc4d19e51d53570be863f6 Mon Sep 17 00:00:00 2001
From: "Erlend E. Aasland" <erlend@python.org>
Date: Thu, 28 Mar 2024 09:40:37 +0100
Subject: [PATCH] gh-116303: Handle disabled test modules in test.support
 helpers (#116482)

Make sure test.support helpers skip iso. failing if test extension
modules are disabled. Also log TEST_MODULES in test.pythoninfo.
---
 Lib/test/pythoninfo.py              |  1 +
 Lib/test/support/__init__.py        | 20 ++++++++++++++++----
 Lib/test/support/bytecode_helper.py | 14 ++++++++++----
 3 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/Lib/test/pythoninfo.py b/Lib/test/pythoninfo.py
index c8bf16dd4d0d49e..0cfd033bb637a7e 100644
--- a/Lib/test/pythoninfo.py
+++ b/Lib/test/pythoninfo.py
@@ -524,6 +524,7 @@ def collect_sysconfig(info_add):
         'Py_GIL_DISABLED',
         'SHELL',
         'SOABI',
+        'TEST_MODULES',
         'abs_builddir',
         'abs_srcdir',
         'prefix',
diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py
index 3d7868768231f50..92e3174407f133d 100644
--- a/Lib/test/support/__init__.py
+++ b/Lib/test/support/__init__.py
@@ -1715,7 +1715,10 @@ def run_in_subinterp(code):
     module is enabled.
     """
     _check_tracemalloc()
-    import _testcapi
+    try:
+        import _testcapi
+    except ImportError:
+        raise unittest.SkipTest("requires _testcapi")
     return _testcapi.run_in_subinterp(code)
 
 
@@ -1725,7 +1728,10 @@ def run_in_subinterp_with_config(code, *, own_gil=None, **config):
     module is enabled.
     """
     _check_tracemalloc()
-    import _testinternalcapi
+    try:
+        import _testinternalcapi
+    except ImportError:
+        raise unittest.SkipTest("requires _testinternalcapi")
     if own_gil is not None:
         assert 'gil' not in config, (own_gil, config)
         config['gil'] = 2 if own_gil else 1
@@ -1887,12 +1893,18 @@ def restore(self):
 
 
 def with_pymalloc():
-    import _testcapi
+    try:
+        import _testcapi
+    except ImportError:
+        raise unittest.SkipTest("requires _testcapi")
     return _testcapi.WITH_PYMALLOC and not Py_GIL_DISABLED
 
 
 def with_mimalloc():
-    import _testcapi
+    try:
+        import _testcapi
+    except ImportError:
+        raise unittest.SkipTest("requires _testcapi")
     return _testcapi.WITH_MIMALLOC
 
 
diff --git a/Lib/test/support/bytecode_helper.py b/Lib/test/support/bytecode_helper.py
index a4845065a5322e3..7a0e884ccc122ae 100644
--- a/Lib/test/support/bytecode_helper.py
+++ b/Lib/test/support/bytecode_helper.py
@@ -3,7 +3,10 @@
 import unittest
 import dis
 import io
-from _testinternalcapi import compiler_codegen, optimize_cfg, assemble_code_object
+try:
+    import _testinternalcapi
+except ImportError:
+    _testinternalcapi = None
 
 _UNSPECIFIED = object()
 
@@ -133,23 +136,26 @@ def complete_insts_info(self, insts):
         return res
 
 
+@unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
 class CodegenTestCase(CompilationStepTestCase):
 
     def generate_code(self, ast):
-        insts, _ = compiler_codegen(ast, "my_file.py", 0)
+        insts, _ = _testinternalcapi.compiler_codegen(ast, "my_file.py", 0)
         return insts
 
 
+@unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
 class CfgOptimizationTestCase(CompilationStepTestCase):
 
     def get_optimized(self, insts, consts, nlocals=0):
         insts = self.normalize_insts(insts)
         insts = self.complete_insts_info(insts)
-        insts = optimize_cfg(insts, consts, nlocals)
+        insts = _testinternalcapi.optimize_cfg(insts, consts, nlocals)
         return insts, consts
 
+@unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
 class AssemblerTestCase(CompilationStepTestCase):
 
     def get_code_object(self, filename, insts, metadata):
-        co = assemble_code_object(filename, insts, metadata)
+        co = _testinternalcapi.assemble_code_object(filename, insts, metadata)
         return co