Add test helper logic to allow for additional commands in lazy import tests

Author: DavidCEllis

Lib/test/support/import_helper.py

@@ -325,7 +325,7 @@ def ready_to_import(name=None, source=""):
                 sys.modules.pop(name, None)
 
 
-def ensure_lazy_imports(imported_module, modules_to_block):
+def ensure_lazy_imports(imported_module, modules_to_block, *, additional_code=None):
     """Test that when imported_module is imported, none of the modules in
     modules_to_block are imported as a side effect."""
     modules_to_block = frozenset(modules_to_block)
@@ -343,6 +343,16 @@ def ensure_lazy_imports(imported_module, modules_to_block):
             raise AssertionError(f'unexpectedly imported after importing {imported_module}: {{after}}')
         """
     )
+    if additional_code:
+        script += additional_code
+        script += textwrap.dedent(
+            f"""
+            if unexpected := modules_to_block & sys.modules.keys():
+                after = ", ".join(unexpected)
+                raise AssertionError(f'unexpectedly imported after additional code: {{after}}')
+            """
+        )
+
     from .script_helper import assert_python_ok
     assert_python_ok("-S", "-c", script)