diff --git a/esphome/espidf/framework.py b/esphome/espidf/framework.py index c6d9ef13f1..badd174421 100644 --- a/esphome/espidf/framework.py +++ b/esphome/espidf/framework.py @@ -1,6 +1,7 @@ """ESP-IDF framework tools for ESPHome.""" from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from ctypes.util import find_library import json import logging @@ -19,6 +20,7 @@ from esphome.build_helpers.ccache import ( from esphome.build_helpers.tools_cache import tools_cache_path from esphome.core import Version from esphome.framework_helpers import ( + BatchDownloadProgress, PathType, archive_extract_all, create_venv, @@ -682,6 +684,12 @@ def _patch_tools_json_demote_unused_tools(framework_path: Path) -> None: ) +# Tool archives are large (tens to hundreds of MB) and served by GitHub / +# dl.espressif.com; a few streams at once saturate most links without +# hammering the host. Smaller than external_files' 8: those are tiny files. +_PREFETCH_WORKERS = 4 + + def _prefetch_idf_tool_archives( framework_path: Path, targets_str: str, @@ -694,10 +702,10 @@ def _prefetch_idf_tool_archives( which makes large archives effectively impossible to fetch on unstable connections (#17703). This asks the framework's idf_tools (via ``get_tool_downloads.py``) which archives the coming install needs, then - downloads each into ``/dist`` with - ``download_with_resume``. The installer then finds the verified archives - already in place ("file ... is already downloaded") and never touches the - network. + downloads them into ``/dist`` with + ``download_with_resume``, a few at a time under one combined progress + bar. The installer then finds the verified archives already in place + ("file ... is already downloaded") and never touches the network. Strictly best-effort: any failure here just logs and returns, leaving ``idf_tools.py install`` to download whatever is missing exactly as @@ -724,21 +732,51 @@ def _prefetch_idf_tool_archives( for entry in json.loads(stdout) if not (dist_path / entry["dest"]).is_file() ] - for index, entry in enumerate(entries, start=1): - _LOGGER.info( - "Downloading %s (%d/%d) ...", entry["name"], index, len(entries) - ) + if not entries: + return + _LOGGER.info( + "Downloading %d ESP-IDF tool archive(s): %s", + len(entries), + ", ".join(entry["name"] for entry in entries), + ) + # tools.json always carries sizes; should one be missing the combined + # bar could not be trusted, so show no bar at all (per-file bars from + # several threads would interleave) rather than a wrong one. + sizes = [entry["size"] for entry in entries] + progress = BatchDownloadProgress( + "Downloading ESP-IDF tools", sum(sizes) if all(sizes) else 0 + ) + # Reported after the bar is done so the warnings do not land on + # its row; list.append is atomic under the GIL. + failures: list[tuple[str, Exception]] = [] + + def _download(entry: dict) -> None: + tracker = progress.tracker() try: download_with_resume( entry["url"], dist_path / entry["dest"], sha256=entry["sha256"], size=entry["size"], + progress=tracker, ) except Exception as e: # noqa: BLE001 # pylint: disable=broad-exception-caught # Keep prefetching the remaining archives; the installer # will retry this one itself (without resume). - _LOGGER.warning("Could not prefetch %s: %s", entry["name"], e) + tracker(0) + failures.append((entry["name"], e)) + + ex = ThreadPoolExecutor(max_workers=min(_PREFETCH_WORKERS, len(entries))) + try: + for future in [ex.submit(_download, entry) for entry in entries]: + future.result() + finally: + # On Ctrl-C drop the queued archives instead of downloading them + # all before the process can exit; in-flight ones still finish. + ex.shutdown(wait=True, cancel_futures=True) + progress.done() + for name, e in failures: + _LOGGER.warning("Could not prefetch %s: %s", name, e) except Exception as e: # noqa: BLE001 # pylint: disable=broad-exception-caught # The installer downloads anything missing itself; never let the # prefetch become a new way for the install to fail. diff --git a/esphome/framework_helpers.py b/esphome/framework_helpers.py index 1a92195e7d..fc9346bed9 100644 --- a/esphome/framework_helpers.py +++ b/esphome/framework_helpers.py @@ -1,7 +1,7 @@ """Generic toolchain installation helpers shared across framework implementations.""" -from collections.abc import Iterable -from contextlib import ExitStack, contextmanager +from collections.abc import Callable, Iterable +from contextlib import ExitStack import hashlib import io import json @@ -24,20 +24,6 @@ PathType = str | os.PathLike _LOGGER = logging.getLogger(__name__) -# Concurrent downloads would interleave their progress bars; a worker thread -# suppresses its bar for the download it runs. -_PROGRESS_LOCAL = threading.local() - - -@contextmanager -def suppress_download_progress(): - """Silence the per-download progress bar in the current thread.""" - _PROGRESS_LOCAL.disabled = True - try: - yield - finally: - _PROGRESS_LOCAL.disabled = False - # Attempts per mirror URL before falling through to the next mirror; only # mid-stream drops retry (resuming when the server gave a validator), @@ -735,7 +721,11 @@ def _response_validator(resp: "requests.Response") -> str | None: def _stream_response_to_file( - resp: "requests.Response", f: IO[bytes], offset: int, size: int | None = None + resp: "requests.Response", + f: IO[bytes], + offset: int, + size: int | None = None, + progress: Callable[[int], None] | None = None, ) -> None: """Stream an open ``_open_ranged`` response body into ``f`` at ``offset``. @@ -743,25 +733,72 @@ def _stream_response_to_file( (effective offset 0) discards the stale bytes. ``offset`` also seeds the progress bar so a resumed download shows overall progress. ``size`` is the known full file size; when None it is derived from the response's - content-length, and without either there is no progress bar. + content-length, and without either there is no progress bar. With + ``progress`` set, no bar is drawn here; the callback gets the absolute + byte count, seeded with ``offset`` and then after each chunk. """ f.seek(offset) f.truncate(offset) total_size = size or offset + _content_length(resp) downloaded = offset - progress = ( - ProgressBar("Downloading") - if total_size > 0 and not getattr(_PROGRESS_LOCAL, "disabled", False) - else None - ) + own_bar: ProgressBar | None = None + if progress is None: + own_bar = ProgressBar("Downloading") if total_size > 0 else None + progress = ( + (lambda done: own_bar.update(done / total_size)) + if own_bar + else (lambda _: None) + ) + progress(downloaded) for chunk in resp.iter_content(chunk_size=256 * 1024): if chunk: f.write(chunk) downloaded += len(chunk) - if progress is not None: - progress.update(downloaded / total_size) - if progress is not None: - progress.update(1) + progress(downloaded) + if own_bar is not None: + own_bar.update(1) + + +class BatchDownloadProgress: + """One progress bar across several concurrent ``download_with_resume`` calls. + + Each ``tracker()`` is a ``progress`` callback for one download; it reports + that file's absolute byte count and the bar shows the sum over ``total``. + The lock also serialises the bar's stderr writes, so worker threads never + interleave frames. With an unknown ``total`` (0) nothing is drawn. Call + ``done()`` once every download has finished (or failed) so a bar that + never reached 100% still ends its line before the next log message. + """ + + def __init__(self, header: str, total: int) -> None: + self._bar = ProgressBar(header) if total > 0 else None + self._total = total + self._sum = 0 + self._lock = threading.Lock() + + def tracker(self) -> Callable[[int], None]: + last = 0 + + def update(done: int) -> None: + nonlocal last + if self._bar is None: + return + with self._lock: + self._sum += done - last + last = done + self._bar.update(min(self._sum / self._total, 1)) + + return update + + def done(self) -> None: + # Nothing to end unless a frame was drawn and it was not the final + # one (update(1) already emitted its own newline). + if ( + self._bar is not None + and self._bar.last_progress is not None + and self._bar.last_progress != 100 + ): + self._bar.done() def download_with_resume( @@ -774,6 +811,7 @@ def download_with_resume( attempts: int = 5, timeout: int = 30, retry_connect_errors: bool = True, + progress: Callable[[int], None] | None = None, ) -> None: """Download ``url`` to ``dest``, resuming partial downloads. @@ -796,6 +834,12 @@ def download_with_resume( of consuming attempts — for callers with their own fallback, like ``download_from_mirrors``. + ``progress``, when given, replaces the built-in progress bar: it is called + with the absolute number of bytes of ``dest`` obtained so far (including + a resumed prefix, and the final size once the file is verified), so a + caller running several downloads at once can draw one combined bar (see + ``BatchDownloadProgress``). + Raises EsphomeError when all attempts are exhausted. """ # Imported lazily: requests is a heavy import (~85ms) and is only needed @@ -819,6 +863,8 @@ def download_with_resume( if dest.is_file() and (sha256 is not None or size is not None): try: _verify_file(dest, sha256, size) + if progress is not None: + progress(size if size is not None else dest.stat().st_size) return except EsphomeError: dest.unlink() @@ -864,7 +910,7 @@ def download_with_resume( # Recorded so a later run can prove an If-Range # resume of this part file safe. _write_download_meta(meta, url, validator, expected_total) - _stream_response_to_file(resp, f, offset, size) + _stream_response_to_file(resp, f, offset, size, progress) # else: a previous run already wrote every byte (or more) but # was killed before the rename below. Skip the network entirely # — a Range request past EOF would draw HTTP 416 — and let @@ -873,6 +919,10 @@ def download_with_resume( expected_size = size if size is not None else expected_total _verify_file(part, sha256, expected_size or None) + if progress is not None: + # Also credits a part file an earlier run completed without + # streaming anything this time. + progress(expected_size or part.stat().st_size) if not expected_size and sha256 is None: # No sha, no size, and the server sent no usable # content-length: nothing can prove the download complete @@ -975,6 +1025,7 @@ def _try_mirrors_once( f: IO[bytes] | None, timeout: int, failures: list[tuple[str, Exception]], + progress: Callable[[int], None] | None = None, ) -> str | None: """Single pass over the resolved mirror ``urls``, one try per URL. @@ -1003,6 +1054,7 @@ def _try_mirrors_once( # next mirror immediately; only mid-stream drops # retry-with-resume on the same URL. retry_connect_errors=False, + progress=progress, ) return url except (requests.RequestException, OSError, EsphomeError) as e: @@ -1044,7 +1096,7 @@ def _try_mirrors_once( if offset == 0: validator = _response_validator(resp) expected_total = _content_length(resp) - _stream_response_to_file(resp, f, offset) + _stream_response_to_file(resp, f, offset, progress=progress) if expected_total and f.tell() != expected_total: raise EsphomeError( @@ -1093,6 +1145,7 @@ def download_from_mirrors( substitutions: dict[str, str], target: io.RawIOBase | IO[bytes] | PathType, timeout: int = 30, + progress: Callable[[int], None] | None = None, ) -> str: """ Download file from multiple mirrors with substitution support. @@ -1102,6 +1155,8 @@ def download_from_mirrors( substitutions: Dictionary of substitutions to apply to URLs target: Target file path or file-like object timeout: Download timeout in seconds + progress: Passed through to the download (see ``download_with_resume``); + replaces the built-in per-file bar Returns: The source URL. @@ -1166,7 +1221,9 @@ def download_from_mirrors( for sweep in range(1, _MIRROR_SWEEP_ATTEMPTS + 1): sweep_failures: list[tuple[str, Exception]] = [] if ( - url := _try_mirrors_once(urls, path_target, f, timeout, sweep_failures) + url := _try_mirrors_once( + urls, path_target, f, timeout, sweep_failures, progress + ) ) is not None: return url failures.extend(sweep_failures) diff --git a/esphome/platformio/library.py b/esphome/platformio/library.py index 097a07b8d2..b1765288fc 100644 --- a/esphome/platformio/library.py +++ b/esphome/platformio/library.py @@ -32,10 +32,10 @@ from urllib.request import url2pathname from esphome import git from esphome.core import CORE, EsphomeError, Library from esphome.framework_helpers import ( + BatchDownloadProgress, archive_extract_all, download_from_mirrors, rmdir, - suppress_download_progress, ) _LOGGER = logging.getLogger(__name__) @@ -86,7 +86,12 @@ ESPHOME_DATA_EXTRA_CMAKE_KEY = "EXTRA_CMAKE" class Source: def download( - self, dir_suffix: str, force: bool = False, salt: str = "", namespace: str = "" + self, + dir_suffix: str, + force: bool = False, + salt: str = "", + namespace: str = "", + progress: Callable[[int], None] | None = None, ) -> Path: raise NotImplementedError @@ -104,7 +109,12 @@ class URLSource(Source): self.url = url def download( - self, dir_suffix: str, force: bool = False, salt: str = "", namespace: str = "" + self, + dir_suffix: str, + force: bool = False, + salt: str = "", + namespace: str = "", + progress: Callable[[int], None] | None = None, ) -> Path: # Namespace the cache per backend (e.g. pio_components/idf, .../zephyr) so # the build files each backend writes into the library dir can't collide. @@ -127,10 +137,12 @@ class URLSource(Source): # Download in temporary file with tempfile.NamedTemporaryFile() as tmp: - _LOGGER.info("Downloading %s ...", self.url) + if progress is None: + # A batch caller draws one combined bar and logs the list + _LOGGER.info("Downloading %s ...", self.url) _LOGGER.debug("Location: %s", path) - download_from_mirrors([self.url], {}, tmp.file) + download_from_mirrors([self.url], {}, tmp.file, progress=progress) _LOGGER.debug("Extracting archive to %s ...", path) archive_extract_all(tmp.file, path) @@ -147,7 +159,12 @@ class GitSource(Source): self.ref = ref def download( - self, dir_suffix: str, force: bool = False, salt: str = "", namespace: str = "" + self, + dir_suffix: str, + force: bool = False, + salt: str = "", + namespace: str = "", + progress: Callable[[int], None] | None = None, ) -> Path: domain = DOMAIN if namespace: @@ -182,7 +199,12 @@ class LocalSource(Source): self.local_path = path def download( - self, dir_suffix: str, force: bool = False, salt: str = "", namespace: str = "" + self, + dir_suffix: str, + force: bool = False, + salt: str = "", + namespace: str = "", + progress: Callable[[int], None] | None = None, ) -> Path: src = Path(self.local_path) if not src.is_dir(): @@ -275,7 +297,13 @@ class ConvertedLibrary: def get_require_name(self): return self.get_sanitized_name().replace("/", "__") - def download(self, force: bool = False, salt: str = "", namespace: str = ""): + def download( + self, + force: bool = False, + salt: str = "", + namespace: str = "", + progress: Callable[[int], None] | None = None, + ): """Fetch the library into the shared cache and record its ``path``. The cache directory is named after the sanitized library name; backends @@ -284,7 +312,11 @@ class ConvertedLibrary: ``get_require_name``). ``namespace`` keeps each backend's cache separate. """ self.path = self.source.download( - self.get_sanitized_name(), force=force, salt=salt, namespace=namespace + self.get_sanitized_name(), + force=force, + salt=salt, + namespace=namespace, + progress=progress, ) self.source_path = self.source.source_root(self.path) @@ -899,6 +931,21 @@ def _warn_unsatisfied_versionless( _DOWNLOAD_WORKERS = 4 +def _content_lengths(urls: list[str]) -> list[int]: + """Content-Length per URL via HEAD requests; 0 for any that fail.""" + import requests + + def head(url: str) -> int: + try: + resp = requests.head(url, timeout=10, allow_redirects=True) + return int(resp.headers.get("content-length", 0)) + except Exception: # noqa: BLE001 # pylint: disable=broad-exception-caught + return 0 + + with ThreadPoolExecutor(max_workers=min(_DOWNLOAD_WORKERS, len(urls))) as ex: + return list(ex.map(head, urls)) + + def _prefetch_wave( wave: list[tuple[str, ConvertedLibrary]], salt: str, namespace: str ) -> None: @@ -920,17 +967,33 @@ def _prefetch_wave( components.append(component) if len(components) < 2: return + _LOGGER.info( + "Downloading %d libraries: %s", + len(components), + ", ".join(c.name for c in components), + ) + # One combined bar over the batch; sizes come from HEAD requests so the + # bar can be trusted (no sizes -> no bar, per BatchDownloadProgress) + sizes = _content_lengths([c.source.url for c in components]) + progress = BatchDownloadProgress( + "Downloading libraries", sum(sizes) if all(sizes) else 0 + ) def _fetch(component: ConvertedLibrary) -> None: + tracker = progress.tracker() try: - with suppress_download_progress(): - component.download(salt=salt, namespace=namespace) + component.download(salt=salt, namespace=namespace, progress=tracker) except Exception: # noqa: BLE001 # pylint: disable=broad-exception-caught # The sequential call below retries and reports the failure - pass + tracker(0) - with ThreadPoolExecutor(max_workers=min(_DOWNLOAD_WORKERS, len(components))) as ex: - list(ex.map(_fetch, components)) + try: + with ThreadPoolExecutor( + max_workers=min(_DOWNLOAD_WORKERS, len(components)) + ) as ex: + list(ex.map(_fetch, components)) + finally: + progress.done() def convert_libraries( diff --git a/tests/unit_tests/test_espidf_component.py b/tests/unit_tests/test_espidf_component.py index e2884454e5..c5f889a37a 100644 --- a/tests/unit_tests/test_espidf_component.py +++ b/tests/unit_tests/test_espidf_component.py @@ -1066,7 +1066,7 @@ def test_idf_component_download_passes_salt() -> None: c.download(force=True, salt="abcd1234", namespace="idf") source.download.assert_called_once_with( - "owner/name", force=True, salt="abcd1234", namespace="idf" + "owner/name", force=True, salt="abcd1234", namespace="idf", progress=None ) assert c.path == Path("/converted/owner/name") diff --git a/tests/unit_tests/test_espidf_framework.py b/tests/unit_tests/test_espidf_framework.py index d66f1fc7db..3a50d0b3be 100644 --- a/tests/unit_tests/test_espidf_framework.py +++ b/tests/unit_tests/test_espidf_framework.py @@ -2,6 +2,7 @@ # pylint: disable=protected-access +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager import importlib.util import io @@ -14,7 +15,7 @@ import subprocess import sys import tarfile from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -895,16 +896,72 @@ def test_prefetch_downloads_each_archive_with_resume(tmp_path: Path) -> None: ), patch("esphome.espidf.framework.download_with_resume") as download, patch("esphome.espidf.framework.get_system_python_path", return_value="python"), + patch("esphome.espidf.framework.BatchDownloadProgress") as progress_cls, ): _prefetch_idf_tool_archives(tmp_path, "esp32", ["required"], None) dist = get_idf_tools_path() / "dist" - assert download.call_count == 2 - assert download.call_args_list[0][0] == ( - "https://example.com/cmake.tar.gz", - dist / "cmake-3.30.2.tar.gz", - ) - assert download.call_args_list[0][1] == {"sha256": "ab" * 32, "size": 123} + # Archives download concurrently, so the call order is not fixed. + calls = {call[0]: call[1] for call in download.call_args_list} + assert set(calls) == { + ("https://example.com/cmake.tar.gz", dist / "cmake-3.30.2.tar.gz"), + ("https://example.com/ninja.zip", dist / "ninja.zip"), + } + kwargs = calls[("https://example.com/cmake.tar.gz", dist / "cmake-3.30.2.tar.gz")] + assert kwargs["sha256"] == "ab" * 32 + assert kwargs["size"] == 123 + # every archive reports into the one combined progress bar + progress_cls.assert_called_once_with("Downloading ESP-IDF tools", 123 + 45) + tracker = progress_cls.return_value.tracker.return_value + assert all(kw["progress"] is tracker for kw in calls.values()) + + +def test_prefetch_downloads_archives_concurrently(tmp_path: Path) -> None: + """More than one archive fans out over a bounded thread pool.""" + entries = [ + { + "name": f"tool{i}@1", + "url": f"https://example.com/tool{i}.tar.gz", + "size": 10, + "sha256": "ab" * 32, + "dest": f"tool{i}.tar.gz", + } + for i in range(6) + ] + with ( + patch( + "esphome.espidf.framework.run_command", + return_value=(True, json.dumps(entries), ""), + ), + patch("esphome.espidf.framework.download_with_resume") as download, + patch("esphome.espidf.framework.get_system_python_path", return_value="python"), + patch( + "esphome.espidf.framework.ThreadPoolExecutor", wraps=ThreadPoolExecutor + ) as pool, + ): + _prefetch_idf_tool_archives(tmp_path, "esp32", ["required"], None) + + pool.assert_called_once_with(max_workers=4) + assert download.call_count == 6 + + +def test_prefetch_single_archive_uses_one_worker(tmp_path: Path) -> None: + entries = json.loads(_PREFETCH_JSON)[:1] + with ( + patch( + "esphome.espidf.framework.run_command", + return_value=(True, json.dumps(entries), ""), + ), + patch("esphome.espidf.framework.download_with_resume") as download, + patch("esphome.espidf.framework.get_system_python_path", return_value="python"), + patch( + "esphome.espidf.framework.ThreadPoolExecutor", wraps=ThreadPoolExecutor + ) as pool, + ): + _prefetch_idf_tool_archives(tmp_path, "esp32", ["required"], None) + + pool.assert_called_once_with(max_workers=1) + assert download.call_count == 1 def test_prefetch_skips_already_downloaded_archives(tmp_path: Path) -> None: @@ -964,6 +1021,11 @@ def test_prefetch_one_failed_archive_does_not_stop_the_rest( ) -> None: """A single archive failing its download must not abort the prefetch of the remaining archives.""" + + def _fail_cmake_download(url: str, *args, **kwargs) -> None: + if "cmake" in url: + raise OSError("network down") + with ( patch( "esphome.espidf.framework.run_command", @@ -971,7 +1033,7 @@ def test_prefetch_one_failed_archive_does_not_stop_the_rest( ), patch( "esphome.espidf.framework.download_with_resume", - side_effect=[OSError("network down"), None], + side_effect=_fail_cmake_download, ) as download, patch("esphome.espidf.framework.get_system_python_path", return_value="python"), ): @@ -981,6 +1043,29 @@ def test_prefetch_one_failed_archive_does_not_stop_the_rest( assert "Could not prefetch cmake@3.30.2" in caplog.text +def test_prefetch_finishes_progress_bar_and_cancels_queue(tmp_path: Path) -> None: + """The batch bar is closed out after the pool, and the pool is shut down + with cancel_futures so Ctrl-C does not drain every queued archive.""" + with ( + patch( + "esphome.espidf.framework.run_command", + return_value=(True, _PREFETCH_JSON, ""), + ), + patch("esphome.espidf.framework.download_with_resume"), + patch("esphome.espidf.framework.get_system_python_path", return_value="python"), + patch("esphome.espidf.framework.BatchDownloadProgress") as progress_cls, + patch( + "esphome.espidf.framework.ThreadPoolExecutor", wraps=ThreadPoolExecutor + ) as pool_cls, + ): + pool = MagicMock(wraps=ThreadPoolExecutor(max_workers=2)) + pool_cls.return_value = pool + _prefetch_idf_tool_archives(tmp_path, "esp32", ["required"], None) + + pool.shutdown.assert_called_once_with(wait=True, cancel_futures=True) + progress_cls.return_value.done.assert_called_once_with() + + def test_prefetch_passes_targets_and_tools_to_script(tmp_path: Path) -> None: with ( patch( diff --git a/tests/unit_tests/test_framework_helpers.py b/tests/unit_tests/test_framework_helpers.py index f7f50b79a4..6be152144a 100644 --- a/tests/unit_tests/test_framework_helpers.py +++ b/tests/unit_tests/test_framework_helpers.py @@ -21,6 +21,7 @@ import requests as req from esphome import framework_helpers from esphome.core import EsphomeError from esphome.framework_helpers import ( + BatchDownloadProgress, _7z_extract_all, _detect_archive_root, _is_transient_download_error, @@ -1112,6 +1113,108 @@ class TestDownloadWithResume: assert mock_get.call_args[1]["headers"] == {} assert dest.read_bytes() == b"data" + def test_progress_callback_reports_absolute_bytes(self, tmp_path: Path) -> None: + """With a callback no bar is drawn; the callback sees the running + byte count of this file, then its final verified size.""" + dest = tmp_path / "tool.tar.gz" + resp = _mock_response(b"") + resp.headers = {"content-length": "7"} + resp.iter_content.return_value = [b"1234", b"567"] + seen: list[int] = [] + with ( + patch("requests.get", return_value=resp), + patch("esphome.framework_helpers.ProgressBar") as bar, + ): + download_with_resume( + "https://example.com/t", dest, size=7, progress=seen.append + ) + assert seen == [0, 4, 7, 7] + bar.assert_not_called() + + def test_progress_callback_seeds_with_resume_offset(self, tmp_path: Path) -> None: + dest = tmp_path / "tool.tar.gz" + (tmp_path / "tool.tar.gz.part").write_bytes(b"12345") + good = hashlib.sha256(b"12345678").hexdigest() + seen: list[int] = [] + with patch("requests.get", return_value=_resumed_response(b"678")): + download_with_resume( + "https://example.com/t", dest, sha256=good, size=8, progress=seen.append + ) + assert seen[0] == 5 + assert seen[-1] == 8 + + def test_progress_callback_credits_already_complete_download( + self, tmp_path: Path + ) -> None: + """A verified dest from an earlier run still counts toward the batch.""" + dest = tmp_path / "tool.tar.gz" + dest.write_bytes(b"12345678") + seen: list[int] = [] + with patch("requests.get") as mock_get: + download_with_resume( + "https://example.com/t", dest, size=8, progress=seen.append + ) + mock_get.assert_not_called() + assert seen == [8] + + +class TestBatchDownloadProgress: + def test_sums_trackers_into_one_bar(self) -> None: + with patch("esphome.framework_helpers.ProgressBar") as bar_cls: + progress = BatchDownloadProgress("Downloading", 100) + a = progress.tracker() + b = progress.tracker() + a(10) + b(20) + a(30) + a(0) # a restart from zero takes that file's bytes back out + bar_cls.assert_called_once_with("Downloading") + updates = [c[0][0] for c in bar_cls.return_value.update.call_args_list] + assert updates == [0.1, 0.3, 0.5, 0.2] + + def test_clamps_at_one(self) -> None: + """Sizes are advisory; an over-delivering server never pushes past 100%.""" + with patch("esphome.framework_helpers.ProgressBar") as bar_cls: + progress = BatchDownloadProgress("Downloading", 10) + progress.tracker()(25) + assert bar_cls.return_value.update.call_args[0][0] == 1 + + def test_unknown_total_draws_nothing(self) -> None: + with patch("esphome.framework_helpers.ProgressBar") as bar_cls: + progress = BatchDownloadProgress("Downloading", 0) + progress.tracker()(5) + progress.done() + bar_cls.assert_not_called() + + def test_done_ends_an_unfinished_bar(self) -> None: + """A batch that stops short of 100% (a failed archive) still ends its + line so the next log message starts on a fresh row.""" + stream = io.StringIO() + stream.isatty = lambda: True # type: ignore[method-assign] + with patch("esphome.helpers.sys.stderr", stream): + progress = BatchDownloadProgress("Downloading", 10) + progress.tracker()(5) + progress.done() + assert stream.getvalue().endswith("50% \n") + + def test_done_before_any_frame_writes_nothing(self) -> None: + """A batch aborted before any tracker fired must not emit a stray + newline for a bar that was never drawn.""" + stream = io.StringIO() + stream.isatty = lambda: True # type: ignore[method-assign] + with patch("esphome.helpers.sys.stderr", stream): + BatchDownloadProgress("Downloading", 10).done() + assert stream.getvalue() == "" + + def test_done_after_full_bar_adds_nothing(self) -> None: + stream = io.StringIO() + stream.isatty = lambda: True # type: ignore[method-assign] + with patch("esphome.helpers.sys.stderr", stream): + progress = BatchDownloadProgress("Downloading", 10) + progress.tracker()(10) + progress.done() + assert stream.getvalue().endswith("100% Done...\r\n") + class TestDownloadFromMirrors: def test_success_returns_url_and_writes_content(self, tmp_path: Path) -> None: @@ -2125,21 +2228,3 @@ def test_strip_win_long_path_prefix( r"""``\\?\`` and ``\\?\UNC\`` prefixes are stripped only on win32.""" with patch("esphome.framework_helpers.sys.platform", platform): assert framework_helpers.strip_win_long_path_prefix(input_path) == expected - - -def test_suppress_download_progress_is_thread_local() -> None: - """The bar suppression only affects the thread that entered the context.""" - import threading - - from esphome import framework_helpers as fh - - seen: list[bool] = [] - with fh.suppress_download_progress(): - assert getattr(fh._PROGRESS_LOCAL, "disabled", False) is True - thread = threading.Thread( - target=lambda: seen.append(getattr(fh._PROGRESS_LOCAL, "disabled", False)) - ) - thread.start() - thread.join() - assert seen == [False] - assert getattr(fh._PROGRESS_LOCAL, "disabled", False) is False diff --git a/tests/unit_tests/test_platformio_library.py b/tests/unit_tests/test_platformio_library.py index ee73384383..35a5ff804a 100644 --- a/tests/unit_tests/test_platformio_library.py +++ b/tests/unit_tests/test_platformio_library.py @@ -4,9 +4,11 @@ Covers the shared download/parse/resolve/dependency-walk paths in ``esphome.platformio.library`` directly (the ESP-IDF and Zephyr backends are exercised in their own test modules).""" +from contextlib import contextmanager import json import logging from pathlib import Path +from types import SimpleNamespace import pytest @@ -157,11 +159,26 @@ def test_localsource_download_returns_empty_build_dir(setup_core: Path) -> None: assert plain != out +@contextmanager +def caplog_at_info(): + records: list[logging.LogRecord] = [] + handler = logging.Handler() + handler.emit = records.append + logger = logging.getLogger("esphome.platformio.library") + logger.addHandler(handler) + try: + yield records + finally: + logger.removeHandler(handler) + + def test_urlsource_download_extracts_then_reuses_marker(setup_core, monkeypatch): monkeypatch.setattr(lib, "rmdir", lambda path, msg="": None) dl_calls: list[list[str]] = [] monkeypatch.setattr( - lib, "download_from_mirrors", lambda urls, headers, f: dl_calls.append(urls) + lib, + "download_from_mirrors", + lambda urls, headers, f, progress=None: dl_calls.append(urls), ) def fake_extract(fileobj, path): @@ -180,6 +197,12 @@ def test_urlsource_download_extracts_then_reuses_marker(setup_core, monkeypatch) assert out2 == out assert len(dl_calls) == 1 + # A batch caller passes a tracker and owns the messaging; no per-file INFO + with caplog_at_info() as records: + src.download("mylib-batch", progress=lambda done: None) + assert len(dl_calls) == 2 + assert not [r for r in records if "Downloading" in r.message] + def test_resolve_registry_version_raises_without_pkg_file(monkeypatch): registry = lib._make_registry_client() @@ -220,7 +243,7 @@ def _patch_registry_resolve(monkeypatch: pytest.MonkeyPatch) -> None: def _patch_download_with_manifests(monkeypatch, tmp_path, manifests, *, properties=()): """Fake ConvertedLibrary.download to materialize canned manifests on disk.""" - def fake_download(self, force=False, salt="", namespace=""): + def fake_download(self, force=False, salt="", namespace="", progress=None): self.path = tmp_path / self.get_require_name() self.path.mkdir(parents=True, exist_ok=True) if self.name in properties: @@ -299,7 +322,11 @@ def _patch_download_without_manifest( calls: list[bool] = [] def fake_download( - self: ConvertedLibrary, force: bool = False, salt: str = "", namespace: str = "" + self: ConvertedLibrary, + force: bool = False, + salt: str = "", + namespace: str = "", + progress=None, ) -> None: calls.append(force) self.path = tmp_path / self.get_require_name() @@ -610,8 +637,10 @@ def test_prefetch_wave_downloads_registry_archives_in_parallel( git/local sources and failures are left to the sequential call.""" calls: list[str] = [] - def fake_download(self, force=False, salt="", namespace=""): + def fake_download(self, force=False, salt="", namespace="", progress=None): calls.append(self.source.url) + if progress is not None: + progress(0) if "boom" in self.source.url: raise RuntimeError("boom") @@ -633,6 +662,22 @@ def test_prefetch_wave_downloads_registry_archives_in_parallel( ] +def test_content_lengths_head_requests(monkeypatch: pytest.MonkeyPatch) -> None: + """Sizes come from HEAD Content-Length; a failing HEAD reads as 0 so + the combined bar is skipped rather than wrong.""" + import requests + + def fake_head(url, timeout, allow_redirects): + if "bad" in url: + raise requests.ConnectionError("down") + return SimpleNamespace(headers={"content-length": "123"}) + + monkeypatch.setattr( + lib.requests if hasattr(lib, "requests") else requests, "head", fake_head + ) + assert lib._content_lengths(["https://x/a", "https://x/bad"]) == [123, 0] + + def test_prefetch_wave_single_archive_skips_the_pool( monkeypatch: pytest.MonkeyPatch, ) -> None: