From de6991144e1a9b7fef89fc82242c48953fae18de Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@home-assistant.io>
Date: Tue, 13 Oct 2026 11:16:05 -1000
Subject: [PATCH] Time only the component import and put the original modules
 back afterwards

---
 .../python/test_component_import_bench.py     | 36 ++++++++++++++-----
 1 file changed, 27 insertions(+), 9 deletions(-)

diff --git a/tests/benchmarks/python/test_component_import_bench.py b/tests/benchmarks/python/test_component_import_bench.py
index a6d0a8ed0f1..0defa3b753e 100644
--- a/tests/benchmarks/python/test_component_import_bench.py
+++ b/tests/benchmarks/python/test_component_import_bench.py
@@ -6,10 +6,11 @@ integrations (mqtt, web_server, zigbee) at module scope drags their whole
 import chains in for every user, so the cost of a fresh import is what these
 benchmarks pin down.
 
-The interpreter is shared across iterations, so each run first drops the
-esphome modules from ``sys.modules`` and imports the component again from
-scratch; the import machinery and the module bodies run every time, as they
-do in a new process.
+The core of esphome is imported once up front and stays loaded; each run
+drops only the component modules from ``sys.modules`` and imports the
+component again, so the import machinery and the module bodies of the
+component and everything it pulls in run every time, as they do in a new
+process, while a change to the core moves no number here.
 
 Skipped when ``pytest-codspeed`` isn't installed so the regular unit-test
 suite keeps working unchanged.
@@ -17,6 +18,7 @@ suite keeps working unchanged.
 
 from __future__ import annotations
 
+from collections.abc import Generator
 import importlib
 import sys
 
@@ -24,19 +26,35 @@ import pytest
 
 pytest.importorskip("pytest_codspeed")
 
+# Load the core before the first case so no case pays for it.
+import esphome.codegen  # noqa: E402, F401
+import esphome.config_validation  # noqa: E402, F401
+
 COMPONENTS = ("binary_sensor", "sensor", "logger")
+COMPONENT_PREFIX = "esphome.components."
 
 
-def _forget_esphome_modules() -> None:
-    for name in [n for n in sys.modules if n == "esphome" or n.startswith("esphome.")]:
-        del sys.modules[name]
+@pytest.fixture
+def restore_modules() -> Generator[None]:
+    """Put the module objects that were loaded before the case back afterwards.
+
+    The case replaces every component module with a fresh copy; anything
+    that imported the originals keeps using them, so they must be the ones
+    left in ``sys.modules`` for later tests.
+    """
+    saved = dict(sys.modules)
+    yield
+    sys.modules.clear()
+    sys.modules.update(saved)
 
 
 def _import_component(name: str) -> object:
-    _forget_esphome_modules()
-    return importlib.import_module(f"esphome.components.{name}")
+    for module in [m for m in sys.modules if m.startswith(COMPONENT_PREFIX)]:
+        del sys.modules[module]
+    return importlib.import_module(COMPONENT_PREFIX + name)
 
 
+@pytest.mark.usefixtures("restore_modules")
 @pytest.mark.parametrize("component", COMPONENTS)
 def test_import_component(benchmark, component: str) -> None:
     benchmark(_import_component, component)