Address review: patch the prefetch subprocess, scope the in-process override to our client, isolate tests

esphome/platformio/library.py

@@ -576,17 +576,21 @@ def _make_registry_client() -> Any:
     elsewhere, not by the PlatformIO registry.
     """
     from platformio.package.manager._registry import PackageManagerRegistryMixin
-
-    from esphome.platformio.runner import patch_registry_private_packages
-
-    # Otherwise every lookup sleeps ~500 ms in PlatformIO's account probe
-    patch_registry_private_packages()
+    from platformio.registry.client import RegistryClient
 
     class _Registry(PackageManagerRegistryMixin):
         def __init__(self) -> None:
             self._registry_client = None
             self.pkg_type = "library"
 
+        def get_registry_client_instance(self) -> RegistryClient:
+            if self._registry_client is None:
+                self._registry_client = RegistryClient()
+                # Skip PlatformIO's account probe: it sleeps ~500 ms per lookup
+                # (see runner.patch_registry_private_packages)
+                self._registry_client.allowed_private_packages = lambda: False
+            return self._registry_client
+
         @staticmethod
         def is_system_compatible(value: Any, custom_system: Any = None) -> bool:
             return True

esphome/platformio/prefetch.py

@@ -922,6 +922,7 @@ def main(argv: list[str]) -> int:
     """Subprocess entry point: ``prefetch <build_dir> <env_name>``."""
     from esphome.core import CORE
     from esphome.log import setup_log
+    from esphome.platformio.runner import patch_registry_private_packages
 
     signal.signal(signal.SIGTERM, _sigterm)
     raw_level = os.environ.get("ESPHOME_PREFETCH_LOG_LEVEL")
@@ -950,6 +951,7 @@ def main(argv: list[str]) -> int:
         _LOGGER.warning("prefetch usage: <build_dir> <env_name>")
         return 2
     build_dir, env = argv
+    patch_registry_private_packages()
     try:
         _prefetch(Path(build_dir), env)
     except KeyboardInterrupt:

tests/unit_tests/test_platformio_library.py

@@ -225,25 +225,24 @@ def test_resolve_registry_version_raises_without_pkg_file(monkeypatch):
 
 
 def test_make_registry_client_skips_private_package_probe(monkeypatch):
-    """Our client never calls PlatformIO's account probe."""
+    """Our client never calls PlatformIO's account probe, and only ours."""
     from platformio.account.client import AccountClient
     from platformio.registry.client import RegistryClient
 
-    # Restore PlatformIO's own probe after the test
-    monkeypatch.setattr(
-        RegistryClient,
-        "allowed_private_packages",
-        RegistryClient.__dict__["allowed_private_packages"],
-    )
+    pio_probe = RegistryClient.__dict__["allowed_private_packages"]
 
     def fail(*_args, **_kwargs):
         raise AssertionError("account probe must not run")
 
     monkeypatch.setattr(AccountClient, "get_account_info", fail)
 
-    client = lib._make_registry_client().get_registry_client_instance()
+    registry = lib._make_registry_client()
+    client = registry.get_registry_client_instance()
 
     assert client.allowed_private_packages() is False
+    assert registry.get_registry_client_instance() is client
+    # Instance override only; the class keeps PlatformIO's probe in this process
+    assert RegistryClient.__dict__["allowed_private_packages"] is pio_probe
 
 
 def _patch_registry_resolve(monkeypatch: pytest.MonkeyPatch) -> None:

tests/unit_tests/test_platformio_prefetch.py

@@ -1151,6 +1151,19 @@ def test_main_runs_prefetch(tmp_path: Path) -> None:
     mock_prefetch.assert_called_once_with(tmp_path, "testenv")
 
 
+def test_main_skips_private_package_probe(tmp_path: Path) -> None:
+    """The registry probe patch is applied before any package manager runs."""
+    with (
+        patch.object(pf, "_prefetch") as mock_prefetch,
+        patch(
+            "esphome.platformio.runner.patch_registry_private_packages"
+        ) as mock_patch,
+    ):
+        assert pf.main([str(tmp_path), "testenv"]) == 0
+    mock_patch.assert_called_once_with()
+    mock_prefetch.assert_called_once()
+
+
 def test_main_bad_argv_is_a_distinct_exit(
     caplog: pytest.LogCaptureFixture,
 ) -> None:

tests/unit_tests/test_platformio_runner.py

@@ -7,6 +7,7 @@ import io
 import sys
 from types import ModuleType
 
+from platformio.registry.client import RegistryClient
 import pytest
 
 from esphome.platformio import runner
@@ -108,25 +109,26 @@ def test_main_applies_registry_private_packages_patch(
     assert order == ["patch", "pio"]
 
 
-def _restore_registry_probe(monkeypatch: pytest.MonkeyPatch) -> None:
-    """Restore PlatformIO's own probe after the test."""
-    from platformio.registry.client import RegistryClient
-
-    monkeypatch.setattr(
-        RegistryClient,
-        "allowed_private_packages",
-        RegistryClient.__dict__["allowed_private_packages"],
-    )
+# Snapshot PlatformIO's own probe at import, before any test can patch it
+_PIO_PROBE = RegistryClient.__dict__["allowed_private_packages"]
 
 
+@pytest.fixture
+def _restore_registry_probe():
+    """Start from and return to PlatformIO's own probe."""
+    RegistryClient.allowed_private_packages = _PIO_PROBE  # type: ignore[method-assign]
+    yield
+    RegistryClient.allowed_private_packages = _PIO_PROBE  # type: ignore[method-assign]
+
+
+@pytest.mark.usefixtures("_restore_registry_probe")
 def test_patch_registry_private_packages_skips_account_probe(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     """Answers False without touching the account client."""
     from platformio.account.client import AccountClient
-    from platformio.registry.client import RegistryClient
 
-    _restore_registry_probe(monkeypatch)
+    assert RegistryClient.__dict__["allowed_private_packages"] is _PIO_PROBE
 
     def fail(*_args, **_kwargs):
         raise AssertionError("account probe must not run")
@@ -139,12 +141,9 @@ def test_patch_registry_private_packages_skips_account_probe(
     assert RegistryClient().allowed_private_packages() is False
 
 
-def test_patch_registry_private_packages_is_idempotent(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    from platformio.registry.client import RegistryClient
-
-    _restore_registry_probe(monkeypatch)
+@pytest.mark.usefixtures("_restore_registry_probe")
+def test_patch_registry_private_packages_is_idempotent() -> None:
+    assert RegistryClient.__dict__["allowed_private_packages"] is _PIO_PROBE
 
     runner.patch_registry_private_packages()
     patched = RegistryClient.allowed_private_packages