diff --git a/esphome/espidf/framework.py b/esphome/espidf/framework.py index 85e9c88171..95d168c85c 100644 --- a/esphome/espidf/framework.py +++ b/esphome/espidf/framework.py @@ -747,27 +747,38 @@ def _prefetch_idf_tool_archives( len(entries), ", ".join(entry["name"] for entry in entries), ) + # tools.json always carries sizes; should one be missing the bar + # could not be trusted, so draw none rather than a wrong one. + sizes = [entry["size"] for entry in entries] progress = BatchDownloadProgress( - "Downloading ESP-IDF tools", - sum(entry["size"] or 0 for entry in entries), + "Downloading ESP-IDF tools", sum(sizes) if all(sizes) else 0 ) 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=progress.tracker(), + 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). + tracker(0) _LOGGER.warning("Could not prefetch %s: %s", entry["name"], e) - with ThreadPoolExecutor(max_workers=min(_PREFETCH_WORKERS, len(entries))) as ex: - list(ex.map(_download, entries)) + 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() 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 76ae6b801f..5247619ab9 100644 --- a/esphome/framework_helpers.py +++ b/esphome/framework_helpers.py @@ -742,7 +742,9 @@ class BatchDownloadProgress: 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. + 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: @@ -765,6 +767,10 @@ class BatchDownloadProgress: return update + def done(self) -> None: + if self._bar is not None and self._bar.last_progress != 100: + self._bar.done() + def download_with_resume( url: str, @@ -829,7 +835,7 @@ def download_with_resume( try: _verify_file(dest, sha256, size) if progress is not None: - progress(dest.stat().st_size) + progress(size if size is not None else dest.stat().st_size) return except EsphomeError: dest.unlink() @@ -887,7 +893,7 @@ def download_with_resume( if progress is not None: # Also credits a part file an earlier run completed without # streaming anything this time. - progress(part.stat().st_size) + 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 diff --git a/tests/unit_tests/test_espidf_framework.py b/tests/unit_tests/test_espidf_framework.py index af14b35d0d..f06190fee7 100644 --- a/tests/unit_tests/test_espidf_framework.py +++ b/tests/unit_tests/test_espidf_framework.py @@ -15,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 @@ -1043,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 ffcb79155d..96da38e606 100644 --- a/tests/unit_tests/test_framework_helpers.py +++ b/tests/unit_tests/test_framework_helpers.py @@ -1181,9 +1181,31 @@ class TestBatchDownloadProgress: def test_unknown_total_draws_nothing(self) -> None: with patch("esphome.framework_helpers.ProgressBar") as bar_cls: - BatchDownloadProgress("Downloading", 0).tracker()(5) + 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_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: