Prefer the landed archive while waiting, cap credited bytes, and give the registry wait a deadline

tests/unit_tests/conftest.py

@@ -143,14 +143,17 @@ def mock_get_component() -> Generator[Mock, None, None]:
 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).
+    Each poll writes the next chunk to ``part`` (or runs it, for a callable)
+    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]
+        part: Path,
+        chunks: list[bytes | Callable[[], None]],
+        land: Callable[[], None],
     ) -> Callable[..., None]:
         polls = iter(chunks)
 
@@ -161,8 +164,11 @@ def held_lock() -> Callable[..., Callable[..., None]]:
                 part.unlink(missing_ok=True)
                 land()
                 return
-            part.parent.mkdir(parents=True, exist_ok=True)
-            part.write_bytes(chunk)
+            if callable(chunk):
+                chunk()
+            else:
+                part.parent.mkdir(parents=True, exist_ok=True)
+                part.write_bytes(chunk)
             raise Timeout("held")
 
         return acquire

tests/unit_tests/test_framework_helpers.py

@@ -2356,7 +2356,7 @@ def test_discard_partial_download_logs_undeletable(
 
 
 def test_downloaded_bytes_reports_what_is_on_disk(tmp_path: Path) -> None:
-    """Landed file: size; part file: its bytes, capped at size; nothing: 0."""
+    """Part file first, then the landed file, both capped at size; else 0."""
     dest = tmp_path / "archive"
     assert framework_helpers.downloaded_bytes(dest, 4) == 0
     part = tmp_path / "archive.part"
@@ -2364,7 +2364,9 @@ def test_downloaded_bytes_reports_what_is_on_disk(tmp_path: Path) -> None:
     assert framework_helpers.downloaded_bytes(dest, 4) == 2
     part.write_bytes(b"abcdef")
     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
+    dest.write_bytes(b"abc")
+    assert framework_helpers.downloaded_bytes(dest, 4) == 3
+    assert framework_helpers.downloaded_bytes(dest) == 3
+    dest.write_bytes(b"abcdef")
+    assert framework_helpers.downloaded_bytes(dest, 4) == 4

tests/unit_tests/test_platformio_prefetch.py

@@ -466,7 +466,7 @@ def test_lock_deadline_leaves_download_to_the_holder(
     with (
         patch("esphome.framework_helpers.download_with_resume") as mock_download,
         patch("filelock.FileLock.acquire", side_effect=Timeout("held")),
-        patch.object(pf, "_DOWNLOAD_LOCK_TIMEOUT", 0),
+        patch("esphome.framework_helpers.DOWNLOAD_LOCK_TIMEOUT", 0),
     ):
         pf._uri_fetch_job(MagicMock(), "https://x/a.zip", dl_path, 4)(ticks.append)
     mock_download.assert_not_called()
@@ -524,6 +524,26 @@ def test_lock_wait_reports_the_holders_progress(
     assert caplog.text.count("Waiting for another process downloading archive") == 1
 
 
+def test_uri_lock_wait_prefers_the_landed_archive(tmp_path: Path, held_lock) -> None:
+    """Between the holder's promotion rename and its release the staging
+    part is gone; the landed cache file is credited instead of 0."""
+    dl_path = tmp_path / "archive"
+    ticks: list[int] = []
+    acquire = held_lock(
+        tmp_path / "archive.prefetch.part",
+        [b"ab", lambda: dl_path.write_bytes(b"abcd")],
+        lambda: None,
+    )
+    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, 4]
+
+
 def test_registry_lock_deadline_skips_registration(tmp_path: Path) -> None:
     """A registry job that lost the download race to another process
     must not stamp a nonexistent archive into pio's usage.db."""
@@ -532,7 +552,7 @@ def test_registry_lock_deadline_skips_registration(tmp_path: Path) -> None:
     with (
         patch("esphome.framework_helpers.download_with_resume") as mock_download,
         patch("filelock.FileLock.acquire", side_effect=Timeout("held")),
-        patch.object(pf, "_DOWNLOAD_LOCK_TIMEOUT", 0),
+        patch("esphome.framework_helpers.DOWNLOAD_LOCK_TIMEOUT", 0),
     ):
         pf._registry_fetch_job(manager, "https://x/a.tar.gz", dl_path, "ab" * 32, 4)(
             lambda done: None

tests/unit_tests/test_platformio_registry.py

@@ -8,6 +8,7 @@ import os
 from pathlib import Path
 from unittest.mock import MagicMock, patch
 
+from filelock import Timeout
 import pytest
 
 from esphome.core import EsphomeError
@@ -564,9 +565,16 @@ def test_prefetch_packages_waits_with_the_holders_progress(
     dest = tmp_path / "a"
     dest.mkdir()
     ticks: list[int] = []
+    part = tmp_path / "dl" / "a-1.0.part"
+
+    def installed_and_pruned() -> None:
+        # install_package touches the marker, then unlinks the archive
+        (dest / ".esphome_extracted").touch()
+        part.unlink()
+
     acquire = held_lock(
-        tmp_path / "dl" / "a-1.0.part",
-        [b"abc"],
+        part,
+        [lambda: None, b"abc", installed_and_pruned],
         (dest / ".esphome_extracted").touch,
     )
 
@@ -588,10 +596,30 @@ def test_prefetch_packages_waits_with_the_holders_progress(
             [("a", "1.0", dest, []), ("b", "2.0", tmp_path / "b", [])],
             tmp_path / "dl",
         )
-    assert ticks == [3, 10]
+    assert ticks == [0, 3, 10, 10]
     mock_download.assert_called_once()
 
 
+def test_prefetch_packages_leaves_a_long_held_lock_to_its_holder(
+    tmp_path: Path,
+) -> None:
+    """Past the deadline the worker skips; install_package waits on the same
+    lock later and verifies whatever the holder produced."""
+    with (
+        patch("filelock.FileLock.acquire", side_effect=Timeout("held")),
+        patch("esphome.framework_helpers.DOWNLOAD_LOCK_TIMEOUT", 0),
+        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", tmp_path / "a", []), ("b", "2.0", tmp_path / "b", [])],
+            tmp_path / "dl",
+        )
+    mock_download.assert_not_called()
+
+
 def test_already_installed_probe(tmp_path: Path) -> None:
     """Both arms of the marker probe the prefetch worker keys on."""
     dest = tmp_path / "pkg"