From f626d73a2b40fadf78e18a4cc4de55bbd6267b7d Mon Sep 17 00:00:00 2001 From: Hommy Date: Sat, 22 Aug 2026 21:54:50 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E8=B5=84=E6=BA=90=E6=96=87?= =?UTF-8?q?=E4=BB=B6=E4=B8=8B=E8=BD=BD=EF=BC=8C=E5=A2=9E=E5=8A=A0=E6=96=AD?= =?UTF-8?q?=E7=82=B9=E7=BB=AD=E4=BC=A0=E7=9A=84=E8=83=BD=E5=8A=9B=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/utils/draft_downloader.py | 246 ++++++++++++++++++----- tests/test_draft_download_resume.py | 292 ++++++++++++++++++++++++++++ 2 files changed, 493 insertions(+), 45 deletions(-) create mode 100644 tests/test_draft_download_resume.py diff --git a/src/utils/draft_downloader.py b/src/utils/draft_downloader.py index 99a1a17..3d9b022 100644 --- a/src/utils/draft_downloader.py +++ b/src/utils/draft_downloader.py @@ -310,6 +310,113 @@ def _http_get(url: str, **kwargs) -> requests.Response: return requests.get(url, headers=headers, **kwargs) +# 草稿文件列表里既有 json 元数据，也有音视频/图片。仅后者需要 HTTP Range 断点续传。 +_MEDIA_RESOURCE_EXTENSIONS = frozenset( + { + ".mp4", + ".mov", + ".avi", + ".mkv", + ".webm", + ".m4v", + ".flv", + ".wmv", + ".ts", + ".mpeg", + ".mpg", + ".3gp", + ".mp3", + ".wav", + ".aac", + ".m4a", + ".flac", + ".ogg", + ".wma", + ".aiff", + ".aif", + ".opus", + ".amr", + ".jpg", + ".jpeg", + ".png", + ".gif", + ".webp", + ".bmp", + ".tiff", + ".tif", + ".heic", + ".heif", + ".ico", + ".svg", + } +) + + +def _extract_extension(url_or_path: str) -> str: + """从 URL 或本地路径取出小写扩展名（忽略 query）。""" + if not url_or_path: + return "" + path = urlparse(url_or_path).path if "://" in url_or_path else url_or_path + return os.path.splitext(path)[1].lower() + + +def _is_media_resource(url_or_path: str) -> bool: + """ + 判断是否为需要断点续传的资源文件（视频/图片/音频）。 + + json、bin、草稿元数据等返回 False，保持「失败即整文件重下」的原有行为。 + """ + return _extract_extension(url_or_path) in _MEDIA_RESOURCE_EXTENSIONS + + +def _local_file_size(path: str) -> int: + """返回本地文件字节数；不存在或无法读取时视为 0。""" + try: + if os.path.isfile(path): + return os.path.getsize(path) + except OSError: + return 0 + return 0 + + +def _resume_request_headers(local_path: str) -> Tuple[Optional[Dict[str, str]], int]: + """ + 按本地半成品大小构造 Range 请求头。 + + 无半成品时返回 (None, 0)，调用方应发普通 GET，请求形态与改造前一致。 + """ + size = _local_file_size(local_path) + if size <= 0: + return None, 0 + return {"Range": f"bytes={size}-"}, size + + +def _is_download_success_status(status_code: int, resume_from: int) -> bool: + """200 始终成功；206 仅在确实携带断点（resume_from>0）时视为续传成功。""" + if status_code == 200: + return True + return status_code == 206 and resume_from > 0 + + +def _write_http_body_to_file( + response: requests.Response, file_path: str, *, append: bool +) -> None: + """ + 将 HTTP 响应体写入本地文件。 + + append=True：断点续传，在已有内容后追加（配合 206）。 + append=False：覆盖写入（JSON 等非资源，或服务端忽略 Range 返回 200 时的整文件重下）。 + """ + parent_dir = os.path.dirname(file_path) + if parent_dir: + os.makedirs(parent_dir, exist_ok=True) + mode = "ab" if append else "wb" + with open(file_path, mode) as out: + for chunk in response.iter_content(chunk_size=8192): + if chunk: + out.write(chunk) + + def _resolve_download_target_path( file_url: str, target_dir: str ) -> Tuple[str, Optional[str]]: @@ -715,16 +822,33 @@ def _download_single_file(file_url: str, target_dir: str) -> None: retry_count = 0 full_file_path, url_draft_id = _resolve_download_target_path(file_url, target_dir) + # 仅视频/图片/音频走 Range 续传；json 等仍每次整文件覆盖下载。 + enable_resume = _is_media_resource(file_url) or _is_media_resource(full_file_path) while retry_count <= max_retries: try: - response = _http_get( - file_url, - timeout=(_REQUEST_CONNECT_TIMEOUT, _REQUEST_READ_TIMEOUT), - stream=True, - ) + extra_headers = None + resume_from = 0 + if enable_resume: + extra_headers, resume_from = _resume_request_headers(full_file_path) + if resume_from > 0: + logger.info( + "Resume media download from byte %s: %s", + resume_from, + file_url, + ) + + get_kwargs = { + "timeout": (_REQUEST_CONNECT_TIMEOUT, _REQUEST_READ_TIMEOUT), + "stream": True, + } + # 无半成品时不传 headers，保持与改造前完全相同的 GET 调用形态。 + if extra_headers: + get_kwargs["headers"] = extra_headers + + response = _http_get(file_url, **get_kwargs) try: - if response.status_code != 200: + if not _is_download_success_status(response.status_code, resume_from): if not _is_retryable_http_status(response.status_code): status = response.status_code logger.error( @@ -763,14 +887,14 @@ def _download_single_file(file_url: str, target_dir: str) -> None: _sleep_transient_http_backoff(retry_count, response) continue - parent_dir = os.path.dirname(full_file_path) - if parent_dir: - os.makedirs(parent_dir, exist_ok=True) - - with open(full_file_path, "wb") as out: - for chunk in response.iter_content(chunk_size=8192): - if chunk: - out.write(chunk) + # 206：从断点追加；200：整文件覆盖（含服务端忽略 Range 的情况）。 + append = response.status_code == 206 and resume_from > 0 + if resume_from > 0 and response.status_code == 200: + logger.info( + "Server ignored Range, re-downloading from start: %s", + file_url, + ) + _write_http_body_to_file(response, full_file_path, append=append) finally: response.close() @@ -1040,15 +1164,31 @@ def _download_remote_material_raising( fallback_ext: str, ) -> str: """下载 URL 素材；失败抛出 DraftDownloadAbort。""" + # 本函数只拉取音视频/图片（含无扩展名的 CDN URL），始终允许断点续传。 + local_path: Optional[str] = None for attempt in range(_MAX_RETRIES + 1): response = None try: - response = _http_get( - file_url, - timeout=(_REQUEST_CONNECT_TIMEOUT, _REQUEST_READ_TIMEOUT), - stream=True, - ) - if response.status_code != 200: + extra_headers = None + resume_from = 0 + if local_path: + extra_headers, resume_from = _resume_request_headers(local_path) + if resume_from > 0: + logger.info( + "Resume remote material from byte %s: %s", + resume_from, + file_url, + ) + + get_kwargs = { + "timeout": (_REQUEST_CONNECT_TIMEOUT, _REQUEST_READ_TIMEOUT), + "stream": True, + } + if extra_headers: + get_kwargs["headers"] = extra_headers + + response = _http_get(file_url, **get_kwargs) + if not _is_download_success_status(response.status_code, resume_from): if not _is_retryable_http_status(response.status_code): status = response.status_code logger.error( @@ -1082,19 +1222,20 @@ def _download_remote_material_raising( response.close() continue - ext = _infer_ext_from_content_type( - response.headers.get("Content-Type"), fallback_ext - ) - filename = _build_material_filename(base_name, ext) - local_path = os.path.join(target_dir, "assets", sub_dir, filename) + if local_path is None: + ext = _infer_ext_from_content_type( + response.headers.get("Content-Type"), fallback_ext + ) + filename = _build_material_filename(base_name, ext) + local_path = os.path.join(target_dir, "assets", sub_dir, filename) - parent_dir = os.path.dirname(local_path) - if parent_dir: - os.makedirs(parent_dir, exist_ok=True) - with open(local_path, "wb") as f: - for chunk in response.iter_content(chunk_size=8192): - if chunk: - f.write(chunk) + append = response.status_code == 206 and resume_from > 0 + if resume_from > 0 and response.status_code == 200: + logger.info( + "Server ignored Range, re-downloading from start: %s", + file_url, + ) + _write_http_body_to_file(response, local_path, append=append) return local_path except DraftDownloadAbort: raise @@ -1152,14 +1293,29 @@ def _download_remote_file(file_url: str, local_path: str) -> bool: def _download_remote_file_raising(file_url: str, local_path: str) -> None: """下载单个 URL 素材；失败抛出 DraftDownloadAbort。""" + enable_resume = _is_media_resource(file_url) or _is_media_resource(local_path) for attempt in range(_MAX_RETRIES + 1): try: - response = _http_get( - file_url, - timeout=(_REQUEST_CONNECT_TIMEOUT, _REQUEST_READ_TIMEOUT), - stream=True, - ) - if response.status_code != 200: + extra_headers = None + resume_from = 0 + if enable_resume: + extra_headers, resume_from = _resume_request_headers(local_path) + if resume_from > 0: + logger.info( + "Resume remote file from byte %s: %s", + resume_from, + file_url, + ) + + get_kwargs = { + "timeout": (_REQUEST_CONNECT_TIMEOUT, _REQUEST_READ_TIMEOUT), + "stream": True, + } + if extra_headers: + get_kwargs["headers"] = extra_headers + + response = _http_get(file_url, **get_kwargs) + if not _is_download_success_status(response.status_code, resume_from): if not _is_retryable_http_status(response.status_code): status = response.status_code logger.error( @@ -1193,13 +1349,13 @@ def _download_remote_file_raising(file_url: str, local_path: str) -> None: response.close() continue - parent_dir = os.path.dirname(local_path) - if parent_dir: - os.makedirs(parent_dir, exist_ok=True) - with open(local_path, "wb") as f: - for chunk in response.iter_content(chunk_size=8192): - if chunk: - f.write(chunk) + append = response.status_code == 206 and resume_from > 0 + if resume_from > 0 and response.status_code == 200: + logger.info( + "Server ignored Range, re-downloading from start: %s", + file_url, + ) + _write_http_body_to_file(response, local_path, append=append) return except DraftDownloadAbort: raise diff --git a/tests/test_draft_download_resume.py b/tests/test_draft_download_resume.py new file mode 100644 index 0000000..cbe2b70 --- /dev/null +++ b/tests/test_draft_download_resume.py @@ -0,0 +1,292 @@ +"""草稿下载：资源文件断点续传；JSON 等非资源仍整文件重下。""" +import os +import tempfile +from unittest.mock import MagicMock, patch + +import pytest +import requests + +import src.utils.draft_downloader as dd + + +@pytest.fixture +def no_sleep(): + with patch.object(dd, "time") as m_time: + m_time.sleep = MagicMock() + yield m_time + + +def _stream_response( + chunks, + status: int = 200, + headers=None, + raise_after=None, +) -> MagicMock: + r = MagicMock() + r.status_code = status + r.headers = headers or {} + r.close = MagicMock() + + def iter_content(chunk_size=8192): + for chunk in chunks: + yield chunk + if raise_after is not None: + raise raise_after + + r.iter_content = iter_content + return r + + +def _range_from_call(call) -> str: + headers = call.kwargs.get("headers") or {} + return headers.get("Range", "") + + +class TestIsMediaResource: + @pytest.mark.parametrize( + "value", + [ + "clip.mp4", + "https://cdn.example.com/a.MP4?token=1", + r"C:\draft\assets\videos\x.mov", + "photo.PNG", + "https://x.test/img.jpg", + "audio.mp3", + "track.wav", + "pic.webp", + ], + ) + def test_media_extensions_are_detected(self, value: str) -> None: + assert dd._is_media_resource(value) is True + + @pytest.mark.parametrize( + "value", + [ + "draft_content.json", + "draft_meta_info.json", + "https://cdn.example.com/app/output/draft/20251204214904ccb1af38/a.bin", + "notes.txt", + "https://x.test/foo.image?sig=1", + "", + ], + ) + def test_non_media_extensions_are_excluded(self, value: str) -> None: + assert dd._is_media_resource(value) is False + + +class TestResumeHelpers: + def test_resume_headers_none_when_file_missing(self, tmp_path) -> None: + headers, resume_from = dd._resume_request_headers(str(tmp_path / "miss.mp4")) + assert headers is None + assert resume_from == 0 + + def test_resume_headers_use_existing_size(self, tmp_path) -> None: + path = tmp_path / "part.mp4" + path.write_bytes(b"hello") + headers, resume_from = dd._resume_request_headers(str(path)) + assert resume_from == 5 + assert headers == {"Range": "bytes=5-"} + + def test_success_status_206_only_when_resuming(self) -> None: + assert dd._is_download_success_status(200, 0) is True + assert dd._is_download_success_status(200, 10) is True + assert dd._is_download_success_status(206, 10) is True + assert dd._is_download_success_status(206, 0) is False + assert dd._is_download_success_status(404, 10) is False + + +class TestDownloadSingleFileResume: + _BASE = "https://capcut.example.com" + _DRAFT = "20251204214904ccb1af38" + _TIMEOUT = (dd._REQUEST_CONNECT_TIMEOUT, dd._REQUEST_READ_TIMEOUT) + _HEADERS = dd._REQUEST_HEADERS + + def _url(self, name: str) -> str: + return f"{self._BASE}/app/output/draft/{self._DRAFT}/{name}" + + def test_first_media_request_has_no_range_header(self, no_sleep) -> None: + """首次下载资源文件不带 Range，GET 形态与改造前一致。""" + file_url = self._url("Resources/clip.mp4") + with tempfile.TemporaryDirectory() as td: + resp = _stream_response([b"ab", b"cd"]) + with patch.object(dd, "requests") as m_req: + m_req.get.return_value = resp + m_req.exceptions = requests.exceptions + assert dd.download_single_file(file_url, td) is True + m_req.get.assert_called_once_with( + file_url, + timeout=self._TIMEOUT, + stream=True, + headers=self._HEADERS, + ) + out = os.path.join(td, "Resources", "clip.mp4") + with open(out, "rb") as f: + assert f.read() == b"abcd" + + def test_media_retry_sends_range_and_appends(self, no_sleep) -> None: + """中途断开后，重试从已写入字节续传，206 响应追加到原文件。""" + file_url = self._url("assets/clip.mp4") + first = _stream_response( + [b"hello"], + raise_after=requests.exceptions.ChunkedEncodingError("truncated"), + ) + second = _stream_response([b"world"], status=206) + + with tempfile.TemporaryDirectory() as td: + with patch.object(dd, "requests") as m_req: + m_req.get.side_effect = [first, second] + m_req.exceptions = requests.exceptions + assert dd.download_single_file(file_url, td) is True + + assert m_req.get.call_count == 2 + assert "Range" not in (m_req.get.call_args_list[0].kwargs.get("headers") or {}) + assert _range_from_call(m_req.get.call_args_list[1]) == "bytes=5-" + out = os.path.join(td, "assets", "clip.mp4") + with open(out, "rb") as f: + assert f.read() == b"helloworld" + + def test_media_retry_overwrites_when_server_ignores_range(self, no_sleep) -> None: + """服务端忽略 Range 返回 200 时，整文件覆盖，避免拼出损坏文件。""" + file_url = self._url("assets/clip.mp4") + first = _stream_response( + [b"hello"], + raise_after=requests.exceptions.ChunkedEncodingError("truncated"), + ) + second = _stream_response([b"FULLFILE"], status=200) + + with tempfile.TemporaryDirectory() as td: + with patch.object(dd, "requests") as m_req: + m_req.get.side_effect = [first, second] + m_req.exceptions = requests.exceptions + assert dd.download_single_file(file_url, td) is True + out = os.path.join(td, "assets", "clip.mp4") + with open(out, "rb") as f: + assert f.read() == b"FULLFILE" + + def test_json_retry_overwrites_without_range(self, no_sleep) -> None: + """JSON 失败重试必须整文件重下，即使本地已有半成品也不发 Range。""" + file_url = self._url("draft_meta_info.json") + first = _stream_response( + [b'{"x":'], + raise_after=requests.exceptions.ChunkedEncodingError("truncated"), + ) + second = _stream_response([b'{"ok": true}']) + + with tempfile.TemporaryDirectory() as td: + with patch.object(dd, "requests") as m_req: + m_req.get.side_effect = [first, second] + m_req.exceptions = requests.exceptions + assert dd.download_single_file(file_url, td) is True + + second_headers = m_req.get.call_args_list[1].kwargs.get("headers") or {} + assert "Range" not in second_headers + out = os.path.join(td, "draft_meta_info.json") + with open(out, "rb") as f: + assert f.read() == b'{"ok": true}' + + def test_bin_retry_overwrites_without_range(self, no_sleep) -> None: + file_url = self._url("assets/x.bin") + first = _stream_response( + [b"123"], + raise_after=requests.exceptions.ReadTimeout("stalled"), + ) + second = _stream_response([b"abc"]) + with tempfile.TemporaryDirectory() as td: + with patch.object(dd, "requests") as m_req: + m_req.get.side_effect = [first, second] + m_req.exceptions = requests.exceptions + assert dd.download_single_file(file_url, td) is True + assert "Range" not in (m_req.get.call_args_list[1].kwargs.get("headers") or {}) + out = os.path.join(td, "assets", "x.bin") + with open(out, "rb") as f: + assert f.read() == b"abc" + + def test_image_resume_same_as_video(self, no_sleep) -> None: + file_url = self._url("Resources/pic.png") + first = _stream_response( + [b"PNG"], + raise_after=requests.exceptions.ChunkedEncodingError("truncated"), + ) + second = _stream_response([b"DATA"], status=206) + with tempfile.TemporaryDirectory() as td: + with patch.object(dd, "requests") as m_req: + m_req.get.side_effect = [first, second] + m_req.exceptions = requests.exceptions + assert dd.download_single_file(file_url, td) is True + assert _range_from_call(m_req.get.call_args_list[1]) == "bytes=3-" + out = os.path.join(td, "Resources", "pic.png") + with open(out, "rb") as f: + assert f.read() == b"PNGDATA" + + +class TestDownloadRemoteMaterialResume: + def test_extensionless_cdn_url_still_resumes(self, no_sleep) -> None: + """无扩展名的图片 CDN URL 也走续传（本函数只下载素材）。""" + url = "https://p3-bot-workflow-sign.byteimg.com/tos-cn-i-mdko3gqilj/foo.png~tplv.image?x=1" + first = _stream_response( + [b"img-"], + headers={"Content-Type": "image/png"}, + raise_after=requests.exceptions.ChunkedEncodingError("truncated"), + ) + second = _stream_response( + [b"rest"], + status=206, + headers={"Content-Type": "image/png"}, + ) + with tempfile.TemporaryDirectory() as td: + with patch.object(dd, "requests") as m_req: + m_req.get.side_effect = [first, second] + m_req.exceptions = requests.exceptions + path = dd._download_remote_material(url, td, "images", "双行", ".mp4") + assert path is not None + assert path.endswith(".png") + with open(path, "rb") as f: + assert f.read() == b"img-rest" + assert _range_from_call(m_req.get.call_args_list[1]) == "bytes=4-" + + def test_first_material_request_has_no_range(self, no_sleep) -> None: + url = "https://cdn.example.com/v?id=1" + resp = _stream_response([b"mp4"], headers={"Content-Type": "video/mp4"}) + with tempfile.TemporaryDirectory() as td: + with patch.object(dd, "requests") as m_req: + m_req.get.return_value = resp + m_req.exceptions = requests.exceptions + path = dd._download_remote_material(url, td, "videos", "clip1", ".bin") + assert path is not None + first_headers = m_req.get.call_args.kwargs.get("headers") or {} + assert "Range" not in first_headers + m_req.get.assert_called_once() + + +class TestDownloadRemoteFileResume: + def test_mp4_retry_appends_on_206(self, no_sleep) -> None: + first = _stream_response( + [b"AAA"], + raise_after=requests.exceptions.ChunkedEncodingError("truncated"), + ) + second = _stream_response([b"BBB"], status=206) + with tempfile.TemporaryDirectory() as td: + out = os.path.join(td, "a.mp4") + with patch.object(dd, "requests") as m_req: + m_req.get.side_effect = [first, second] + m_req.exceptions = requests.exceptions + assert dd._download_remote_file("https://x.test/a.mp4", out) is True + assert _range_from_call(m_req.get.call_args_list[1]) == "bytes=3-" + with open(out, "rb") as f: + assert f.read() == b"AAABBB" + + def test_non_media_remote_file_does_not_resume(self, no_sleep) -> None: + first = _stream_response( + [b"AAA"], + raise_after=requests.exceptions.ChunkedEncodingError("truncated"), + ) + second = _stream_response([b"ZZZ"]) + with tempfile.TemporaryDirectory() as td: + out = os.path.join(td, "a.bin") + with patch.object(dd, "requests") as m_req: + m_req.get.side_effect = [first, second] + m_req.exceptions = requests.exceptions + assert dd._download_remote_file("https://x.test/a.bin", out) is True + assert "Range" not in (m_req.get.call_args_list[1].kwargs.get("headers") or {}) + with open(out, "rb") as f: + assert f.read() == b"ZZZ"