From 160bc6a3e1297c2841d4eb20d3d4d356bd002892 Mon Sep 17 00:00:00 2001 From: Alexander Piskun Date: Thu, 12 Mar 2026 09:41:32 +0000 Subject: [PATCH] fix: handle missing Content-Length header in download_file --- comfy_cli/file_utils.py | 9 +++++++-- tests/test_file_utils_network.py | 18 ++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/comfy_cli/file_utils.py b/comfy_cli/file_utils.py index b58f998e..4bf488dc 100644 --- a/comfy_cli/file_utils.py +++ b/comfy_cli/file_utils.py @@ -69,13 +69,18 @@ def download_file(url: str, local_filepath: pathlib.Path, headers: dict | None = with httpx.stream("GET", url, follow_redirects=True, headers=headers) as response: if response.status_code == 200: - total = int(response.headers["Content-Length"]) + content_length = response.headers.get("Content-Length") + total = int(content_length) if content_length is not None else None + if total is not None: + description = f"Downloading {total // 1024 // 1024} MB" + else: + description = "Downloading..." try: with open(local_filepath, "wb") as f: for data in ui.show_progress( response.iter_bytes(), total, - description=f"Downloading {total // 1024 // 1024} MB", + description=description, ): f.write(data) except KeyboardInterrupt: diff --git a/tests/test_file_utils_network.py b/tests/test_file_utils_network.py index 3ee6055d..7e82e2af 100644 --- a/tests/test_file_utils_network.py +++ b/tests/test_file_utils_network.py @@ -85,6 +85,24 @@ def test_download_file_success(mock_stream, tmp_path): assert test_file.read_bytes() == b"test data" +@patch("httpx.stream") +def test_download_file_success_without_content_length(mock_stream, tmp_path): + """Download should succeed when Content-Length header is missing (e.g. chunked/gzip responses).""" + mock_response = Mock() + mock_response.status_code = 200 + mock_response.headers = {} + mock_response.iter_bytes.return_value = [b"chunk1", b"chunk2"] + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=None) + mock_stream.return_value = mock_response + + test_file = tmp_path / "test.txt" + download_file("http://example.com", test_file) + + assert test_file.exists() + assert test_file.read_bytes() == b"chunk1chunk2" + + @patch("httpx.stream") def test_download_file_failure(mock_stream): mock_response = Mock()