diff --git a/tests/unittests/isolated_import_utils.py b/tests/unittests/isolated_import_utils.py index b77fc316aaa..7c8e7b09bdb 100644 --- a/tests/unittests/isolated_import_utils.py +++ b/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( diff --git a/tests/unittests/test_import_loading.py b/tests/unittests/test_import_loading.py index c4834dfb216..f170b43d9e6 100644 --- a/tests/unittests/test_import_loading.py +++ b/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: