test: measure packages introduced by isolated imports

Author: astrogilda

tests/unittests/isolated_import_utils.py

@@ -56,7 +56,11 @@ def run_isolated(source: str) -> subprocess.CompletedProcess[str]:
 
 
 def loaded_top_level_packages(source: str) -> frozenset[str]:
-  """Returns the third-party top-level packages source leaves imported.
+  """Returns third-party top-level packages introduced by source.
+
+  Interpreter startup hooks can import packages before source runs. Those
+  modules are not costs introduced by source. Snapshot full module names so a
+  newly imported submodule still counts even when its parent was preloaded.
 
   Standard-library modules, private modules and the pseudo-modules the
   interpreter injects carry no install or startup cost of their own, so they
@@ -70,12 +74,14 @@ def loaded_top_level_packages(source: str) -> frozenset[str]:
   result = run_isolated(f"""
 import json
 import sys
+modules_before = set(sys.modules)
 {source}
 
 names = {{
     name.partition('.')[0]
     for name, module in sys.modules.items()
-    if getattr(module, '__spec__', None) is not None
+    if name not in modules_before
+    and getattr(module, '__spec__', None) is not None
     and module.__spec__.origin is not None
 }}
 print(json.dumps(sorted(

tests/unittests/test_import_loading.py

@@ -247,6 +247,43 @@ def test_entry_point_loads_only_allowlisted_packages(statement: str) -> None:
   )
 
 
+@pytest.fixture
+def preloaded_startup_package(tmp_path, monkeypatch):
+  package = tmp_path / 'startup_package'
+  package.mkdir()
+  (package / '__init__.py').write_text('value = 1\n')
+  (package / 'extra.py').write_text('value = 2\n')
+  (tmp_path / 'sitecustomize.py').write_text('import startup_package\n')
+  monkeypatch.setenv('PYTHONPATH', str(tmp_path))
+  return tmp_path
+
+
+def test_import_measurement_excludes_interpreter_startup(
+    preloaded_startup_package,
+):
+  assert (
+      loaded_top_level_packages("assert 'startup_package' in sys.modules")
+      == frozenset()
+  )
+
+
+def test_import_measurement_detects_new_package_after_startup(
+    preloaded_startup_package,
+):
+  (preloaded_startup_package / 'new_package.py').write_text('value = 3\n')
+  assert loaded_top_level_packages(
+      "assert 'startup_package' in sys.modules\nimport new_package"
+  ) == frozenset({'new_package'})
+
+
+def test_import_measurement_detects_new_submodule_of_preloaded_package(
+    preloaded_startup_package,
+):
+  assert loaded_top_level_packages(
+      "assert 'startup_package' in sys.modules\nimport startup_package.extra"
+  ) == frozenset({'startup_package'})
+
+
 def test_constructing_agent_defers_optional_mcp_server_stack():
   """A normal Agent does not import MCP just because its extra is installed."""
   if importlib.util.find_spec('mcp') is None: