Share the lock wait with the ESP-IDF registry prefetch and trim the docstrings

tests/unit_tests/conftest.py

@@ -9,7 +9,7 @@ not be part of a unit test suite.
 
 """
 
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 import os
 from pathlib import Path
 import sys
@@ -137,3 +137,34 @@ def mock_get_component() -> Generator[Mock, None, None]:
     """Mock get_component for config module."""
     with patch("esphome.config.get_component") as mock:
         yield mock
+
+
+@pytest.fixture
+def held_lock() -> Callable[..., Callable[..., None]]:
+    """Factory for a ``FileLock.acquire`` fake held by another downloader.
+
+    Each poll writes the next chunk to ``part`` and raises ``Timeout``; when
+    the chunks run out the part is removed, ``land()`` runs, and the acquire
+    succeeds (also for any later job, so ``land`` must be idempotent).
+    """
+    from filelock import Timeout
+
+    def make(
+        part: Path, chunks: list[bytes], land: Callable[[], None]
+    ) -> Callable[..., None]:
+        polls = iter(chunks)
+
+        def acquire(*args, **kwargs) -> None:
+            try:
+                chunk = next(polls)
+            except StopIteration:
+                part.unlink(missing_ok=True)
+                land()
+                return
+            part.parent.mkdir(parents=True, exist_ok=True)
+            part.write_bytes(chunk)
+            raise Timeout("held")
+
+        return acquire
+
+    return make

tests/unit_tests/test_framework_helpers.py

@@ -2356,8 +2356,7 @@ def test_discard_partial_download_logs_undeletable(
 
 
 def test_downloaded_bytes_reports_what_is_on_disk(tmp_path: Path) -> None:
-    """A landed file counts in full, a part file counts what it holds
-    (never more than the size), and nothing on disk counts zero."""
+    """Landed file: size; part file: its bytes, capped at size; nothing: 0."""
     dest = tmp_path / "archive"
     assert framework_helpers.downloaded_bytes(dest, 4) == 0
     part = tmp_path / "archive.part"
@@ -2367,3 +2366,5 @@ def test_downloaded_bytes_reports_what_is_on_disk(tmp_path: Path) -> None:
     assert framework_helpers.downloaded_bytes(dest, 4) == 4
     dest.write_bytes(b"abcd")
     assert framework_helpers.downloaded_bytes(dest, 4) == 4
+    part.unlink()
+    assert framework_helpers.downloaded_bytes(dest) == 4

tests/unit_tests/test_platformio_prefetch.py

@@ -458,13 +458,10 @@ def test_uri_fetch_job_waits_out_a_briefly_held_lock(tmp_path: Path) -> None:
 def test_lock_deadline_leaves_download_to_the_holder(
     tmp_path: Path, staged: bytes
 ) -> None:
-    """A lock held past the deadline means another process is fetching the
-    same file; skipping cleanly beats a misleading failure warning. The
-    tracker is still polled so a parked worker observes cancellation, and
-    it reports what the holder has staged so far."""
+    """A lock held past the deadline is another process's download; skip
+    cleanly, polling the tracker with what the holder has staged so far."""
     dl_path = tmp_path / "archive"
-    if staged:
-        (tmp_path / "archive.prefetch.part").write_bytes(staged)
+    (tmp_path / "archive.prefetch.part").write_bytes(staged)
     ticks: list[int] = []
     with (
         patch("esphome.framework_helpers.download_with_resume") as mock_download,
@@ -477,67 +474,54 @@ def test_lock_deadline_leaves_download_to_the_holder(
     assert not dl_path.exists()
 
 
+@pytest.mark.parametrize(
+    ("job", "part_name", "chunks", "expected"),
+    [
+        (
+            lambda dl_path: pf._registry_fetch_job(
+                MagicMock(), "https://x/a.tar.gz", dl_path, "ab" * 32, 4
+            ),
+            "archive.part",
+            [b"a", b"abc"],
+            [1, 3, 4],
+        ),
+        (
+            lambda dl_path: pf._uri_fetch_job(
+                MagicMock(), "https://x/a.zip", dl_path, 4
+            ),
+            "archive.prefetch.part",
+            [b"ab"],
+            [2, 4],
+        ),
+    ],
+    ids=["registry", "uri"],
+)
 def test_lock_wait_reports_the_holders_progress(
-    tmp_path: Path, caplog: pytest.LogCaptureFixture
+    tmp_path: Path,
+    caplog: pytest.LogCaptureFixture,
+    held_lock,
+    job,
+    part_name: str,
+    chunks: list[bytes],
+    expected: list[int],
 ) -> None:
-    """While another process holds the lock the bar follows its part file,
-    and the job credits the full size once that process lands the archive
-    (the 0% bar of a build racing the dashboard for the same toolchain)."""
-    manager = MagicMock()
+    """A waiting job reports the holder's part file (the staging one for a
+    URL job), then the full size once the holder lands the archive."""
     dl_path = tmp_path / "archive"
-    part = tmp_path / "archive.part"
     ticks: list[int] = []
-    polls = iter([b"a", b"abc"])
-
-    def acquire(*args, **kwargs):
-        try:
-            part.write_bytes(next(polls))
-        except StopIteration:
-            part.unlink()
-            dl_path.write_bytes(b"abcd")
-            return
-        raise Timeout("held")
-
+    acquire = held_lock(
+        tmp_path / part_name, chunks, lambda: dl_path.write_bytes(b"abcd")
+    )
     with (
         patch("esphome.framework_helpers.download_with_resume") as mock_download,
         patch("filelock.FileLock.acquire", side_effect=acquire),
         patch("filelock.FileLock.release"),
         caplog.at_level(logging.INFO),
     ):
-        pf._registry_fetch_job(manager, "https://x/a.tar.gz", dl_path, "ab" * 32, 4)(
-            ticks.append
-        )
+        job(dl_path)(ticks.append)
     mock_download.assert_not_called()
-    assert ticks == [1, 3, 4]
+    assert ticks == expected
     assert caplog.text.count("Waiting for another process downloading archive") == 1
-    manager.set_download_utime.assert_called_once_with(str(dl_path))
-
-
-def test_uri_lock_wait_reads_the_staging_part(tmp_path: Path) -> None:
-    """A URL job's holder streams into the staging path, so the wait
-    reports that part file, not one beside the cache path."""
-    dl_path = tmp_path / "archive"
-    staging_part = tmp_path / "archive.prefetch.part"
-    ticks: list[int] = []
-    polls = iter([b"ab"])
-
-    def acquire(*args, **kwargs):
-        try:
-            staging_part.write_bytes(next(polls))
-        except StopIteration:
-            staging_part.unlink()
-            dl_path.write_bytes(b"abcd")
-            return
-        raise Timeout("held")
-
-    with (
-        patch("esphome.framework_helpers.download_with_resume") as mock_download,
-        patch("filelock.FileLock.acquire", side_effect=acquire),
-        patch("filelock.FileLock.release"),
-    ):
-        pf._uri_fetch_job(MagicMock(), "https://x/a.zip", dl_path, 4)(ticks.append)
-    mock_download.assert_not_called()
-    assert ticks == [2, 4]
 
 
 def test_registry_lock_deadline_skips_registration(tmp_path: Path) -> None:

tests/unit_tests/test_platformio_registry.py

@@ -540,16 +540,13 @@ def test_prefetch_packages_skips_freshly_installed_dest(tmp_path: Path) -> None:
     dest = tmp_path / "a"
     dest.mkdir()
 
-    from contextlib import contextmanager
-
-    @contextmanager
-    def marker_appears_under_lock(path, **kwargs):
+    def marker_appears_under_lock(*args, **kwargs):
         # Simulates the concurrent build finishing while we waited
         (dest / ".esphome_extracted").touch()
-        yield
 
     with (
-        patch("filelock.FileLock", side_effect=marker_appears_under_lock),
+        patch("filelock.FileLock.acquire", side_effect=marker_appears_under_lock),
+        patch("filelock.FileLock.release"),
         patch.object(registry, "download_with_resume") as mock_download,
         patch.object(
             registry, "registry_download", side_effect=_resolve_for({"a": 10})
@@ -559,6 +556,42 @@ def test_prefetch_packages_skips_freshly_installed_dest(tmp_path: Path) -> None:
     mock_download.assert_not_called()
 
 
+def test_prefetch_packages_waits_with_the_holders_progress(
+    tmp_path: Path, held_lock
+) -> None:
+    """A worker parked on another build's lock reports that build's part
+    file, then the full size once the marker appears."""
+    dest = tmp_path / "a"
+    dest.mkdir()
+    ticks: list[int] = []
+    acquire = held_lock(
+        tmp_path / "dl" / "a-1.0.part",
+        [b"abc"],
+        (dest / ".esphome_extracted").touch,
+    )
+
+    def fake_batch(header, jobs):
+        for _name, _size, fetch in jobs:
+            fetch(ticks.append)
+        return []
+
+    with (
+        patch("filelock.FileLock.acquire", side_effect=acquire),
+        patch("filelock.FileLock.release"),
+        patch.object(registry, "run_batch_downloads", side_effect=fake_batch),
+        patch.object(registry, "download_with_resume") as mock_download,
+        patch.object(
+            registry, "registry_download", side_effect=_resolve_for({"a": 10, "b": 5})
+        ),
+    ):
+        registry.prefetch_packages(
+            [("a", "1.0", dest, []), ("b", "2.0", tmp_path / "b", [])],
+            tmp_path / "dl",
+        )
+    assert ticks == [3, 10]
+    mock_download.assert_called_once()
+
+
 def test_already_installed_probe(tmp_path: Path) -> None:
     """Both arms of the marker probe the prefetch worker keys on."""
     dest = tmp_path / "pkg"