| 1 | import asyncio |
| 2 | import json |
| 3 | from datetime import datetime |
| 4 | from unittest.mock import AsyncMock, MagicMock |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from auth import CookieManager |
| 9 | from config import ConfigLoader |
| 10 | from control import QueueManager, RateLimiter, RetryHandler |
| 11 | from core.api_client import DouyinAPIClient |
| 12 | from core.video_downloader import VideoDownloader |
| 13 | from storage import FileManager |
| 14 | |
| 15 | |
| 16 | class _FakeProgressReporter: |
| 17 | def __init__(self): |
| 18 | self.step_updates = [] |
| 19 | self.item_totals = [] |
| 20 | self.item_events = [] |
| 21 | |
| 22 | def update_step(self, step: str, detail: str = "") -> None: |
| 23 | self.step_updates.append((step, detail)) |
| 24 | |
| 25 | def set_item_total(self, total: int, detail: str = "") -> None: |
| 26 | self.item_totals.append((total, detail)) |
| 27 | |
| 28 | def advance_item(self, status: str, detail: str = "") -> None: |
| 29 | self.item_events.append((status, detail)) |
| 30 | |
| 31 | |
| 32 | def _build_downloader(tmp_path): |
| 33 | config = ConfigLoader() |
| 34 | config.update(path=str(tmp_path)) |
| 35 | |
| 36 | file_manager = FileManager(str(tmp_path)) |
| 37 | cookie_manager = CookieManager(str(tmp_path / ".cookies.json")) |
| 38 | api_client = DouyinAPIClient({}) |
| 39 | |
| 40 | downloader = VideoDownloader( |
| 41 | config, |
| 42 | api_client, |
| 43 | file_manager, |
| 44 | cookie_manager, |
| 45 | database=None, |
| 46 | rate_limiter=RateLimiter(max_per_second=5), |
| 47 | retry_handler=RetryHandler(max_retries=1), |
| 48 | queue_manager=QueueManager(max_workers=1), |
| 49 | ) |
| 50 | |
| 51 | return downloader, api_client |
| 52 | |
| 53 | |
| 54 | @pytest.mark.asyncio |
| 55 | async def test_video_downloader_skip_counts_total(tmp_path, monkeypatch): |
| 56 | downloader, api_client = _build_downloader(tmp_path) |
| 57 | |
| 58 | async def _fake_should_download(self, _): |
| 59 | return False |
| 60 | |
| 61 | downloader._should_download = _fake_should_download.__get__(downloader, VideoDownloader) |
| 62 | |
| 63 | result = await downloader.download({"aweme_id": "123"}) |
| 64 | |
| 65 | assert result.total == 1 |
| 66 | assert result.skipped == 1 |
| 67 | assert result.success == 0 |
| 68 | assert result.failed == 0 |
| 69 | |
| 70 | await api_client.close() |
| 71 | |
| 72 | |
| 73 | @pytest.mark.asyncio |
| 74 | async def test_video_downloader_reports_item_progress(tmp_path, monkeypatch): |
| 75 | downloader, api_client = _build_downloader(tmp_path) |
| 76 | reporter = _FakeProgressReporter() |
| 77 | downloader.progress_reporter = reporter |
| 78 | |
| 79 | async def _fake_should_download(self, _aweme_id): |
| 80 | return True |
| 81 | |
| 82 | async def _fake_get_video_detail(_aweme_id: str): |
| 83 | return {"aweme_id": "123", "author": {"nickname": "tester"}} |
| 84 | |
| 85 | async def _fake_download_aweme(self, _aweme_data): |
| 86 | return True |
| 87 | |
| 88 | downloader._should_download = _fake_should_download.__get__(downloader, VideoDownloader) |
| 89 | monkeypatch.setattr(api_client, "get_video_detail", _fake_get_video_detail) |
| 90 | downloader._download_aweme = _fake_download_aweme.__get__(downloader, VideoDownloader) |
| 91 | |
| 92 | result = await downloader.download({"aweme_id": "123"}) |
| 93 | |
| 94 | assert result.total == 1 |
| 95 | assert result.success == 1 |
| 96 | assert reporter.item_totals == [(1, "单作品下载")] |
| 97 | assert ("下载作品", "单作品资源下载中") in reporter.step_updates |
| 98 | assert reporter.item_events == [("success", "123")] |
| 99 | |
| 100 | await api_client.close() |
| 101 | |
| 102 | |
| 103 | @pytest.mark.asyncio |
| 104 | async def test_video_downloader_downloads_note_video_fallback(tmp_path, monkeypatch): |
| 105 | downloader, api_client = _build_downloader(tmp_path) |
| 106 | downloader.config.update(music=False, cover=False, avatar=False, json=False, folderstyle=True) |
| 107 | |
| 108 | aweme_id = "7646971177114611826" |
| 109 | |
| 110 | async def _fake_should_download(self, _aweme_id): |
| 111 | return True |
| 112 | |
| 113 | async def _fake_get_video_detail(_aweme_id: str): |
| 114 | assert _aweme_id == aweme_id |
| 115 | return { |
| 116 | "aweme_id": aweme_id, |
| 117 | "aweme_type": 68, |
| 118 | "desc": "note 视频作品", |
| 119 | "video": { |
| 120 | "play_addr_h264": { |
| 121 | "url_list": ["https://v3-web.douyinvod.com/note-h264.mp4"] |
| 122 | } |
| 123 | }, |
| 124 | } |
| 125 | |
| 126 | async def _fake_get_session(): |
| 127 | return object() |
| 128 | |
| 129 | saved = [] |
| 130 | |
| 131 | async def _fake_download_with_retry(self, url, save_path, _session, **_kwargs): |
| 132 | saved.append((url, save_path)) |
| 133 | return True |
| 134 | |
| 135 | downloader._should_download = _fake_should_download.__get__(downloader, VideoDownloader) |
| 136 | monkeypatch.setattr(api_client, "get_video_detail", _fake_get_video_detail) |
| 137 | monkeypatch.setattr(api_client, "get_session", _fake_get_session) |
| 138 | downloader._download_with_retry = _fake_download_with_retry.__get__(downloader, VideoDownloader) |
| 139 | |
| 140 | result = await downloader.download({"type": "gallery", "aweme_id": aweme_id}) |
| 141 | |
| 142 | assert result.total == 1 |
| 143 | assert result.success == 1 |
| 144 | assert result.failed == 0 |
| 145 | assert saved[0][0] == "https://v3-web.douyinvod.com/note-h264.mp4" |
| 146 | assert saved[0][1].suffix == ".mp4" |
| 147 | |
| 148 | await api_client.close() |
| 149 | |
| 150 | |
| 151 | @pytest.mark.asyncio |
| 152 | async def test_build_no_watermark_url_signs_with_headers(tmp_path, monkeypatch): |
| 153 | downloader, api_client = _build_downloader(tmp_path) |
| 154 | |
| 155 | signed_url = "https://www.douyin.com/aweme/v1/play/?video_id=1&X-Bogus=signed" |
| 156 | |
| 157 | def _fake_sign(url: str): |
| 158 | return signed_url, "UnitTestAgent/1.0" |
| 159 | |
| 160 | monkeypatch.setattr(api_client, "sign_url", _fake_sign) |
| 161 | |
| 162 | aweme = { |
| 163 | "aweme_id": "1", |
| 164 | "video": { |
| 165 | "play_addr": { |
| 166 | "url_list": ["https://www.douyin.com/aweme/v1/play/?video_id=1&watermark=0"] |
| 167 | } |
| 168 | }, |
| 169 | } |
| 170 | |
| 171 | url, headers = downloader._build_no_watermark_url(aweme) |
| 172 | |
| 173 | assert url == signed_url |
| 174 | assert headers["User-Agent"] == "UnitTestAgent/1.0" |
| 175 | assert headers["Accept"] == "*/*" |
| 176 | assert headers["Referer"].startswith("https://www.douyin.com") |
| 177 | |
| 178 | await api_client.close() |
| 179 | |
| 180 | |
| 181 | @pytest.mark.asyncio |
| 182 | async def test_build_no_watermark_url_avoids_playwm_when_uri_can_be_signed(tmp_path, monkeypatch): |
| 183 | downloader, api_client = _build_downloader(tmp_path) |
| 184 | |
| 185 | signed_url = "https://www.douyin.com/aweme/v1/play/?video_id=clean&watermark=0" |
| 186 | |
| 187 | def _fake_build_signed_path(path, params): |
| 188 | assert path == "/aweme/v1/play/" |
| 189 | assert params["video_id"] == "clean" |
| 190 | assert params["watermark"] == "0" |
| 191 | return signed_url, "UnitTestAgent/2.0" |
| 192 | |
| 193 | monkeypatch.setattr(api_client, "build_signed_path", _fake_build_signed_path) |
| 194 | |
| 195 | aweme = { |
| 196 | "aweme_id": "1", |
| 197 | "video": { |
| 198 | "play_addr": { |
| 199 | "uri": "clean", |
| 200 | "url_list": ["https://v3-web.douyinvod.com/playwm/abc.mp4?watermark=1"], |
| 201 | } |
| 202 | }, |
| 203 | } |
| 204 | |
| 205 | url, headers = downloader._build_no_watermark_url(aweme) |
| 206 | |
| 207 | assert url == signed_url |
| 208 | assert headers["User-Agent"] == "UnitTestAgent/2.0" |
| 209 | |
| 210 | await api_client.close() |
| 211 | |
| 212 | |
| 213 | @pytest.mark.asyncio |
| 214 | async def test_build_no_watermark_url_prefers_signed_uri_when_variant_exists( |
| 215 | tmp_path, monkeypatch |
| 216 | ): |
| 217 | downloader, api_client = _build_downloader(tmp_path) |
| 218 | |
| 219 | signed_url = "https://www.douyin.com/aweme/v1/play/?video_id=clean&watermark=0" |
| 220 | |
| 221 | def _fake_build_signed_path(path, params): |
| 222 | assert path == "/aweme/v1/play/" |
| 223 | assert params["video_id"] == "clean" |
| 224 | return signed_url, "UnitTestAgent/2.1" |
| 225 | |
| 226 | monkeypatch.setattr(api_client, "build_signed_path", _fake_build_signed_path) |
| 227 | |
| 228 | aweme = { |
| 229 | "aweme_id": "1", |
| 230 | "video": { |
| 231 | "play_addr_h264": { |
| 232 | "url_list": ["https://v3-web.douyinvod.com/direct-h264.mp4"] |
| 233 | }, |
| 234 | "play_addr": { |
| 235 | "uri": "clean", |
| 236 | "url_list": ["https://v3-web.douyinvod.com/playwm/abc.mp4?watermark=1"], |
| 237 | }, |
| 238 | }, |
| 239 | } |
| 240 | |
| 241 | url, headers = downloader._build_no_watermark_url(aweme) |
| 242 | |
| 243 | assert url == signed_url |
| 244 | assert headers["User-Agent"] == "UnitTestAgent/2.1" |
| 245 | |
| 246 | await api_client.close() |
| 247 | |
| 248 | |
| 249 | @pytest.mark.asyncio |
| 250 | async def test_should_download_skips_when_aweme_exists_locally(tmp_path): |
| 251 | downloader, api_client = _build_downloader(tmp_path) |
| 252 | aweme_id = "7600223638943468863" |
| 253 | |
| 254 | existing_file = tmp_path / f"2026-02-18_demo_{aweme_id}.mp4" |
| 255 | existing_file.write_bytes(b"1") |
| 256 | |
| 257 | should_download = await downloader._should_download(aweme_id) |
| 258 | assert should_download is False |
| 259 | |
| 260 | await api_client.close() |
| 261 | |
| 262 | |
| 263 | @pytest.mark.asyncio |
| 264 | async def test_download_aweme_assets_uses_publish_date_and_writes_manifest(tmp_path, monkeypatch): |
| 265 | downloader, api_client = _build_downloader(tmp_path) |
| 266 | downloader.config.update(music=False, cover=False, avatar=False, json=False, folderstyle=True) |
| 267 | |
| 268 | async def _fake_get_session(): |
| 269 | return object() |
| 270 | |
| 271 | monkeypatch.setattr(api_client, "get_session", _fake_get_session) |
| 272 | |
| 273 | saved_paths = [] |
| 274 | |
| 275 | async def _fake_download_with_retry(self, _url, save_path, _session, **_kwargs): |
| 276 | saved_paths.append(save_path) |
| 277 | return True |
| 278 | |
| 279 | downloader._download_with_retry = _fake_download_with_retry.__get__(downloader, VideoDownloader) |
| 280 | |
| 281 | aweme_id = "7600224486650121526" |
| 282 | publish_ts = 1707303025 |
| 283 | expected_date_prefix = datetime.fromtimestamp(publish_ts).strftime("%Y-%m-%d") |
| 284 | aweme_data = { |
| 285 | "aweme_id": aweme_id, |
| 286 | "desc": "测试下载日期文件名", |
| 287 | "create_time": publish_ts, |
| 288 | "text_extra": [{"hashtag_name": "测试标签"}], |
| 289 | "video": {"play_addr": {"url_list": ["https://example.com/video.mp4"]}}, |
| 290 | } |
| 291 | |
| 292 | success = await downloader._download_aweme_assets( |
| 293 | aweme_data, author_name="测试作者", mode="post" |
| 294 | ) |
| 295 | |
| 296 | assert success is True |
| 297 | assert len(saved_paths) == 1 |
| 298 | |
| 299 | save_path = saved_paths[0] |
| 300 | assert save_path.name.startswith(f"{expected_date_prefix}_") |
| 301 | assert aweme_id in save_path.name |
| 302 | assert save_path.parent.name.startswith(f"{expected_date_prefix}_") |
| 303 | |
| 304 | manifest_path = tmp_path / "download_manifest.jsonl" |
| 305 | assert manifest_path.exists() |
| 306 | lines = manifest_path.read_text(encoding="utf-8").strip().splitlines() |
| 307 | assert len(lines) == 1 |
| 308 | |
| 309 | manifest_entry = json.loads(lines[0]) |
| 310 | assert manifest_entry["date"] == expected_date_prefix |
| 311 | assert manifest_entry["aweme_id"] == aweme_id |
| 312 | assert manifest_entry["tags"] == ["测试标签"] |
| 313 | assert save_path.name in manifest_entry["file_names"] |
| 314 | |
| 315 | await api_client.close() |
| 316 | |
| 317 | |
| 318 | @pytest.mark.asyncio |
| 319 | async def test_download_aweme_assets_keeps_success_when_transcript_skipped(tmp_path, monkeypatch): |
| 320 | downloader, api_client = _build_downloader(tmp_path) |
| 321 | downloader.config.update( |
| 322 | music=False, |
| 323 | cover=False, |
| 324 | avatar=False, |
| 325 | json=False, |
| 326 | folderstyle=True, |
| 327 | transcript={ |
| 328 | "enabled": True, |
| 329 | "api_key_env": "OPENAI_API_KEY", |
| 330 | "api_key": "", |
| 331 | "output_dir": "", |
| 332 | "response_formats": ["txt", "json"], |
| 333 | }, |
| 334 | ) |
| 335 | |
| 336 | async def _fake_get_session(): |
| 337 | return object() |
| 338 | |
| 339 | monkeypatch.setattr(api_client, "get_session", _fake_get_session) |
| 340 | |
| 341 | async def _fake_download_with_retry(self, _url, _save_path, _session, **_kwargs): |
| 342 | return True |
| 343 | |
| 344 | downloader._download_with_retry = _fake_download_with_retry.__get__(downloader, VideoDownloader) |
| 345 | |
| 346 | aweme_data = { |
| 347 | "aweme_id": "7600224486650121527", |
| 348 | "desc": "转写缺 key 也不应影响下载", |
| 349 | "video": {"play_addr": {"url_list": ["https://example.com/video.mp4"]}}, |
| 350 | } |
| 351 | |
| 352 | success = await downloader._download_aweme_assets( |
| 353 | aweme_data, author_name="测试作者", mode="post" |
| 354 | ) |
| 355 | |
| 356 | assert success is True |
| 357 | |
| 358 | await api_client.close() |
| 359 | |
| 360 | |
| 361 | @pytest.mark.asyncio |
| 362 | async def test_download_aweme_assets_video_writes_cover_avatar_and_json(tmp_path, monkeypatch): |
| 363 | downloader, api_client = _build_downloader(tmp_path) |
| 364 | downloader.config.update( |
| 365 | music=False, |
| 366 | cover=True, |
| 367 | avatar=True, |
| 368 | json=True, |
| 369 | folderstyle=True, |
| 370 | transcript={"enabled": False}, |
| 371 | ) |
| 372 | |
| 373 | async def _fake_get_session(): |
| 374 | return object() |
| 375 | |
| 376 | monkeypatch.setattr(api_client, "get_session", _fake_get_session) |
| 377 | |
| 378 | saved_paths = [] |
| 379 | |
| 380 | async def _fake_download_with_retry(self, _url, save_path, _session, **_kwargs): |
| 381 | saved_paths.append(save_path) |
| 382 | return True |
| 383 | |
| 384 | downloader._download_with_retry = _fake_download_with_retry.__get__(downloader, VideoDownloader) |
| 385 | |
| 386 | aweme_data = { |
| 387 | "aweme_id": "7600224486650121527", |
| 388 | "desc": "附加资源", |
| 389 | "create_time": 1707303025, |
| 390 | "author": { |
| 391 | "nickname": "测试作者", |
| 392 | "avatar_larger": {"url_list": ["https://example.com/avatar.jpg"]}, |
| 393 | }, |
| 394 | "video": { |
| 395 | "play_addr": {"url_list": ["https://example.com/video.mp4"]}, |
| 396 | "cover": {"url_list": ["https://example.com/cover.jpg"]}, |
| 397 | }, |
| 398 | } |
| 399 | |
| 400 | success = await downloader._download_aweme_assets( |
| 401 | aweme_data, author_name="测试作者", mode="post" |
| 402 | ) |
| 403 | |
| 404 | assert success is True |
| 405 | assert any(path.name.endswith(".mp4") for path in saved_paths) |
| 406 | assert any(path.name.endswith("_cover.jpg") for path in saved_paths) |
| 407 | assert any(path.name.endswith("_avatar.jpg") for path in saved_paths) |
| 408 | metadata_files = list(tmp_path.rglob("*_data.json")) |
| 409 | assert len(metadata_files) == 1 |
| 410 | |
| 411 | await api_client.close() |
| 412 | |
| 413 | |
| 414 | @pytest.mark.asyncio |
| 415 | async def test_download_aweme_assets_gallery_downloads_live_photo_videos(tmp_path, monkeypatch): |
| 416 | downloader, api_client = _build_downloader(tmp_path) |
| 417 | downloader.config.update(music=False, cover=False, avatar=False, json=False, folderstyle=True) |
| 418 | |
| 419 | async def _fake_get_session(): |
| 420 | return object() |
| 421 | |
| 422 | monkeypatch.setattr(api_client, "get_session", _fake_get_session) |
| 423 | |
| 424 | saved_paths = [] |
| 425 | |
| 426 | async def _fake_download_with_retry(self, _url, save_path, _session, **_kwargs): |
| 427 | saved_paths.append(save_path) |
| 428 | return True |
| 429 | |
| 430 | downloader._download_with_retry = _fake_download_with_retry.__get__(downloader, VideoDownloader) |
| 431 | |
| 432 | aweme_data = { |
| 433 | "aweme_id": "7600224486650121528", |
| 434 | "desc": "实况图文", |
| 435 | "image_post_info": { |
| 436 | "images": [ |
| 437 | { |
| 438 | "display_image": {"url_list": ["https://example.com/1.webp"]}, |
| 439 | "video": {"play_addr": {"url_list": ["https://example.com/1_live.mp4"]}}, |
| 440 | }, |
| 441 | { |
| 442 | "video": {"play_addr": {"url_list": ["https://example.com/2_live.mp4"]}}, |
| 443 | }, |
| 444 | ] |
| 445 | }, |
| 446 | } |
| 447 | |
| 448 | success = await downloader._download_aweme_assets( |
| 449 | aweme_data, author_name="测试作者", mode="post" |
| 450 | ) |
| 451 | |
| 452 | assert success is True |
| 453 | assert any(path.suffix == ".webp" for path in saved_paths) |
| 454 | assert sum(path.suffix == ".mp4" for path in saved_paths) == 2 |
| 455 | assert any("_live_1.mp4" in path.name for path in saved_paths) |
| 456 | assert any("_live_2.mp4" in path.name for path in saved_paths) |
| 457 | |
| 458 | await api_client.close() |
| 459 | |
| 460 | |
| 461 | @pytest.mark.asyncio |
| 462 | async def test_download_aweme_assets_gallery_preserves_real_image_extensions(tmp_path, monkeypatch): |
| 463 | downloader, api_client = _build_downloader(tmp_path) |
| 464 | downloader.config.update(music=False, cover=False, avatar=False, json=False, folderstyle=True) |
| 465 | |
| 466 | async def _fake_get_session(): |
| 467 | return object() |
| 468 | |
| 469 | monkeypatch.setattr(api_client, "get_session", _fake_get_session) |
| 470 | |
| 471 | saved_paths = [] |
| 472 | |
| 473 | async def _fake_download_with_retry(self, _url, save_path, _session, **_kwargs): |
| 474 | saved_paths.append(save_path) |
| 475 | return True |
| 476 | |
| 477 | downloader._download_with_retry = _fake_download_with_retry.__get__(downloader, VideoDownloader) |
| 478 | |
| 479 | aweme_data = { |
| 480 | "aweme_id": "7600224486650121991", |
| 481 | "desc": "图集后缀归一化", |
| 482 | "image_post_info": { |
| 483 | "images": [ |
| 484 | { |
| 485 | "display_image": { |
| 486 | "url_list": ["https://example.com/gallery_1.png~tplv-obj.image?x=1"] |
| 487 | } |
| 488 | }, |
| 489 | { |
| 490 | "display_image": { |
| 491 | "url_list": ["https://example.com/gallery_2.jpeg~tplv-resize:1080:0.image"] |
| 492 | } |
| 493 | }, |
| 494 | { |
| 495 | "display_image": { |
| 496 | "url_list": ["https://example.com/gallery_3.jpg?from=unit-test"] |
| 497 | } |
| 498 | }, |
| 499 | ] |
| 500 | }, |
| 501 | } |
| 502 | |
| 503 | success = await downloader._download_aweme_assets( |
| 504 | aweme_data, author_name="测试作者", mode="post" |
| 505 | ) |
| 506 | |
| 507 | assert success is True |
| 508 | assert [path.suffix for path in saved_paths] == [".png", ".jpeg", ".jpg"] |
| 509 | |
| 510 | await api_client.close() |
| 511 | |
| 512 | |
| 513 | @pytest.mark.asyncio |
| 514 | async def test_download_aweme_assets_gallery_uses_response_content_type_for_suffix( |
| 515 | tmp_path, monkeypatch |
| 516 | ): |
| 517 | downloader, api_client = _build_downloader(tmp_path) |
| 518 | downloader.config.update(music=False, cover=False, avatar=False, json=False, folderstyle=True) |
| 519 | |
| 520 | content = b"fake png content" |
| 521 | publish_ts = 1707303025 |
| 522 | publish_date = datetime.fromtimestamp(publish_ts).strftime("%Y-%m-%d") |
| 523 | aweme_id = "7600224486650121992" |
| 524 | |
| 525 | mock_response = AsyncMock() |
| 526 | mock_response.status = 200 |
| 527 | mock_response.content_length = len(content) |
| 528 | mock_response.headers = {"Content-Type": "image/png; charset=binary"} |
| 529 | |
| 530 | async def iter_chunked(_size): |
| 531 | yield content |
| 532 | |
| 533 | mock_response.content = MagicMock() |
| 534 | mock_response.content.iter_chunked = iter_chunked |
| 535 | |
| 536 | ctx = AsyncMock() |
| 537 | ctx.__aenter__ = AsyncMock(return_value=mock_response) |
| 538 | ctx.__aexit__ = AsyncMock(return_value=False) |
| 539 | |
| 540 | mock_session = MagicMock() |
| 541 | mock_session.get.return_value = ctx |
| 542 | |
| 543 | async def _fake_get_session(): |
| 544 | return mock_session |
| 545 | |
| 546 | monkeypatch.setattr(api_client, "get_session", _fake_get_session) |
| 547 | |
| 548 | aweme_data = { |
| 549 | "aweme_id": aweme_id, |
| 550 | "desc": "响应头决定后缀", |
| 551 | "create_time": publish_ts, |
| 552 | "image_post_info": { |
| 553 | "images": [{"display_image": {"url_list": ["https://example.com/gallery_1.image?x=1"]}}] |
| 554 | }, |
| 555 | } |
| 556 | |
| 557 | success = await downloader._download_aweme_assets( |
| 558 | aweme_data, author_name="测试作者", mode="post" |
| 559 | ) |
| 560 | |
| 561 | assert success is True |
| 562 | save_dir = tmp_path / "测试作者" / "post" / f"{publish_date}_响应头决定后缀_{aweme_id}" |
| 563 | saved_files = sorted(path.name for path in save_dir.iterdir() if path.is_file()) |
| 564 | assert saved_files == [f"{publish_date}_响应头决定后缀_{aweme_id}_1.png"] |
| 565 | |
| 566 | manifest_path = tmp_path / "download_manifest.jsonl" |
| 567 | lines = manifest_path.read_text(encoding="utf-8").strip().splitlines() |
| 568 | manifest_entry = json.loads(lines[-1]) |
| 569 | assert manifest_entry["file_names"] == saved_files |
| 570 | |
| 571 | await api_client.close() |
| 572 | |
| 573 | |
| 574 | @pytest.mark.asyncio |
| 575 | async def test_download_aweme_assets_gallery_tries_next_image_candidate(tmp_path, monkeypatch): |
| 576 | downloader, api_client = _build_downloader(tmp_path) |
| 577 | downloader.config.update(music=False, cover=False, avatar=False, json=False, folderstyle=True) |
| 578 | |
| 579 | async def _fake_get_session(): |
| 580 | return object() |
| 581 | |
| 582 | monkeypatch.setattr(api_client, "get_session", _fake_get_session) |
| 583 | |
| 584 | attempted_urls = [] |
| 585 | |
| 586 | async def _fake_download_with_retry(self, url, _save_path, _session, **_kwargs): |
| 587 | attempted_urls.append(url) |
| 588 | return url.endswith("good.jpeg") |
| 589 | |
| 590 | downloader._download_with_retry = _fake_download_with_retry.__get__(downloader, VideoDownloader) |
| 591 | |
| 592 | aweme_data = { |
| 593 | "aweme_id": "7600224486650121993", |
| 594 | "desc": "候选图回退", |
| 595 | "image_post_info": { |
| 596 | "images": [ |
| 597 | { |
| 598 | "download_url_list": [ |
| 599 | "https://example.com/bad.jpg", |
| 600 | "https://example.com/good.jpeg", |
| 601 | ], |
| 602 | "url_list": ["https://example.com/preview.webp"], |
| 603 | } |
| 604 | ] |
| 605 | }, |
| 606 | } |
| 607 | |
| 608 | success = await downloader._download_aweme_assets( |
| 609 | aweme_data, author_name="测试作者", mode="post" |
| 610 | ) |
| 611 | |
| 612 | assert success is True |
| 613 | assert attempted_urls == [ |
| 614 | "https://example.com/preview.webp", |
| 615 | "https://example.com/bad.jpg", |
| 616 | "https://example.com/good.jpeg", |
| 617 | ] |
| 618 | |
| 619 | await api_client.close() |
| 620 | |
| 621 | |
| 622 | def test_collect_image_urls_prefers_jpeg_over_webp_companion(tmp_path): |
| 623 | downloader, api_client = _build_downloader(tmp_path) |
| 624 | |
| 625 | aweme_data = { |
| 626 | "aweme_id": "100006", |
| 627 | "images": [ |
| 628 | { |
| 629 | "download_url_list": [ |
| 630 | "https://example.com/image.webp", |
| 631 | "https://example.com/image.jpeg", |
| 632 | ], |
| 633 | }, |
| 634 | ], |
| 635 | } |
| 636 | |
| 637 | urls = downloader._collect_image_urls(aweme_data) |
| 638 | |
| 639 | assert urls == ["https://example.com/image.jpeg"] |
| 640 | |
| 641 | asyncio.run(api_client.close()) |
| 642 | |
| 643 | |
| 644 | @pytest.mark.asyncio |
| 645 | async def test_download_aweme_assets_gallery_succeeds_with_only_live_videos(tmp_path, monkeypatch): |
| 646 | downloader, api_client = _build_downloader(tmp_path) |
| 647 | downloader.config.update(music=False, cover=False, avatar=False, json=False, folderstyle=True) |
| 648 | |
| 649 | async def _fake_get_session(): |
| 650 | return object() |
| 651 | |
| 652 | monkeypatch.setattr(api_client, "get_session", _fake_get_session) |
| 653 | |
| 654 | saved_paths = [] |
| 655 | |
| 656 | async def _fake_download_with_retry(self, _url, save_path, _session, **_kwargs): |
| 657 | saved_paths.append(save_path) |
| 658 | return True |
| 659 | |
| 660 | downloader._download_with_retry = _fake_download_with_retry.__get__(downloader, VideoDownloader) |
| 661 | |
| 662 | aweme_data = { |
| 663 | "aweme_id": "7600224486650121529", |
| 664 | "desc": "仅实况图文", |
| 665 | "image_post_info": { |
| 666 | "images": [ |
| 667 | {"video": {"play_addr": {"url_list": ["https://example.com/only_live_1.mp4"]}}}, |
| 668 | {"video": {"play_addr": {"url_list": ["https://example.com/only_live_2.mp4"]}}}, |
| 669 | ] |
| 670 | }, |
| 671 | } |
| 672 | |
| 673 | success = await downloader._download_aweme_assets( |
| 674 | aweme_data, author_name="测试作者", mode="post" |
| 675 | ) |
| 676 | |
| 677 | assert success is True |
| 678 | assert len(saved_paths) == 2 |
| 679 | assert all(path.suffix == ".mp4" for path in saved_paths) |
| 680 | assert any("_live_1.mp4" in path.name for path in saved_paths) |
| 681 | assert any("_live_2.mp4" in path.name for path in saved_paths) |
| 682 | |
| 683 | await api_client.close() |
| 684 | |
| 685 | |
| 686 | @pytest.mark.asyncio |
| 687 | async def test_download_aweme_assets_gallery_fails_when_live_video_download_fails( |
| 688 | tmp_path, monkeypatch |
| 689 | ): |
| 690 | downloader, api_client = _build_downloader(tmp_path) |
| 691 | downloader.config.update(music=False, cover=False, avatar=False, json=False, folderstyle=True) |
| 692 | |
| 693 | async def _fake_get_session(): |
| 694 | return object() |
| 695 | |
| 696 | monkeypatch.setattr(api_client, "get_session", _fake_get_session) |
| 697 | |
| 698 | saved_paths = [] |
| 699 | |
| 700 | async def _fake_download_with_retry(self, _url, save_path, _session, **_kwargs): |
| 701 | saved_paths.append(save_path) |
| 702 | if save_path.name.endswith("_live_2.mp4"): |
| 703 | return False |
| 704 | return True |
| 705 | |
| 706 | downloader._download_with_retry = _fake_download_with_retry.__get__(downloader, VideoDownloader) |
| 707 | |
| 708 | aweme_data = { |
| 709 | "aweme_id": "7600224486650121530", |
| 710 | "desc": "实况下载失败场景", |
| 711 | "image_post_info": { |
| 712 | "images": [ |
| 713 | { |
| 714 | "display_image": {"url_list": ["https://example.com/ok.webp"]}, |
| 715 | "video": {"play_addr": {"url_list": ["https://example.com/live_ok.mp4"]}}, |
| 716 | }, |
| 717 | {"video": {"play_addr": {"url_list": ["https://example.com/live_fail.mp4"]}}}, |
| 718 | ] |
| 719 | }, |
| 720 | } |
| 721 | |
| 722 | success = await downloader._download_aweme_assets( |
| 723 | aweme_data, author_name="测试作者", mode="post" |
| 724 | ) |
| 725 | |
| 726 | assert success is False |
| 727 | assert any(path.name.endswith(".webp") for path in saved_paths) |
| 728 | assert any(path.name.endswith("_live_1.mp4") for path in saved_paths) |
| 729 | assert any(path.name.endswith("_live_2.mp4") for path in saved_paths) |
| 730 | |
| 731 | await api_client.close() |
| 732 | |
| 733 | |
| 734 | def test_detect_media_type_by_aweme_type(tmp_path): |
| 735 | """aweme_type 2/68/150 should be detected as gallery even without images key.""" |
| 736 | downloader, api_client = _build_downloader(tmp_path) |
| 737 | |
| 738 | for aweme_type in (2, 68, 150): |
| 739 | assert downloader._detect_media_type({"aweme_type": aweme_type}) == "gallery" |
| 740 | |
| 741 | assert downloader._detect_media_type({"aweme_type": 4}) == "video" |
| 742 | assert downloader._detect_media_type({"aweme_type": 0}) == "video" |
| 743 | assert downloader._detect_media_type({}) == "video" |
| 744 | |
| 745 | asyncio.run(api_client.close()) |
| 746 | |
| 747 | |
| 748 | def test_collect_image_urls_old_format_url_list(tmp_path): |
| 749 | """Old format: items have url_list directly.""" |
| 750 | downloader, api_client = _build_downloader(tmp_path) |
| 751 | |
| 752 | aweme_data = { |
| 753 | "aweme_id": "100001", |
| 754 | "images": [ |
| 755 | {"url_list": ["https://example.com/img1.webp"]}, |
| 756 | {"url_list": ["https://example.com/img2.webp"]}, |
| 757 | ], |
| 758 | } |
| 759 | |
| 760 | urls = downloader._collect_image_urls(aweme_data) |
| 761 | assert urls == [ |
| 762 | "https://example.com/img1.webp", |
| 763 | "https://example.com/img2.webp", |
| 764 | ] |
| 765 | |
| 766 | asyncio.run(api_client.close()) |
| 767 | |
| 768 | |
| 769 | def test_collect_image_urls_old_format_prefers_url_list(tmp_path): |
| 770 | """Old format: url_list is the no-watermark image source.""" |
| 771 | downloader, api_client = _build_downloader(tmp_path) |
| 772 | |
| 773 | aweme_data = { |
| 774 | "aweme_id": "100002", |
| 775 | "images": [ |
| 776 | { |
| 777 | "url_list": ["https://example.com/preview1.webp"], |
| 778 | "download_url_list": ["https://example.com/download1.webp"], |
| 779 | }, |
| 780 | ], |
| 781 | } |
| 782 | |
| 783 | urls = downloader._collect_image_urls(aweme_data) |
| 784 | assert urls == ["https://example.com/preview1.webp"] |
| 785 | |
| 786 | asyncio.run(api_client.close()) |
| 787 | |
| 788 | |
| 789 | def test_collect_image_urls_new_format_prefers_display_image(tmp_path): |
| 790 | """New format: display_image is the no-watermark image source.""" |
| 791 | downloader, api_client = _build_downloader(tmp_path) |
| 792 | |
| 793 | aweme_data = { |
| 794 | "aweme_id": "100003", |
| 795 | "image_post_info": { |
| 796 | "images": [ |
| 797 | { |
| 798 | "download_url": {"url_list": ["https://cdn.example.com/download.webp"]}, |
| 799 | "display_image": {"url_list": ["https://cdn.example.com/display.webp"]}, |
| 800 | }, |
| 801 | ] |
| 802 | }, |
| 803 | } |
| 804 | |
| 805 | urls = downloader._collect_image_urls(aweme_data) |
| 806 | assert urls == ["https://cdn.example.com/display.webp"] |
| 807 | |
| 808 | asyncio.run(api_client.close()) |
| 809 | |
| 810 | |
| 811 | def test_collect_image_urls_prefers_aweme_image_url_list_before_display_image(tmp_path): |
| 812 | downloader, api_client = _build_downloader(tmp_path) |
| 813 | |
| 814 | aweme_data = { |
| 815 | "aweme_id": "100003-url-list", |
| 816 | "image_post_info": { |
| 817 | "images": [ |
| 818 | { |
| 819 | "url_list": ["https://cdn.example.com/clean-from-aweme.webp"], |
| 820 | "display_image": { |
| 821 | "url_list": ["https://cdn.example.com/tplv-dy-water-v2/display.webp"] |
| 822 | }, |
| 823 | "download_url_list": ["https://cdn.example.com/tplv-dy-water-v2/download.webp"], |
| 824 | }, |
| 825 | ] |
| 826 | }, |
| 827 | } |
| 828 | |
| 829 | urls = downloader._collect_image_urls(aweme_data) |
| 830 | assert urls == ["https://cdn.example.com/clean-from-aweme.webp"] |
| 831 | |
| 832 | asyncio.run(api_client.close()) |
| 833 | |
| 834 | |
| 835 | def test_collect_image_urls_prefers_non_watermark_gallery_fields(tmp_path): |
| 836 | downloader, api_client = _build_downloader(tmp_path) |
| 837 | |
| 838 | aweme_data = { |
| 839 | "aweme_id": "100004", |
| 840 | "image_post_info": { |
| 841 | "images": [ |
| 842 | { |
| 843 | "display_image": {"url_list": ["https://cdn.example.com/clean-display.webp"]}, |
| 844 | "download_url": { |
| 845 | "url_list": ["https://cdn.example.com/tplv-dy-water-v2/water-download.webp"] |
| 846 | }, |
| 847 | "owner_watermark_image": { |
| 848 | "url_list": ["https://cdn.example.com/owner_watermark_image.webp"] |
| 849 | }, |
| 850 | }, |
| 851 | { |
| 852 | "url_list": ["https://cdn.example.com/clean-top.webp"], |
| 853 | "download_url_list": [ |
| 854 | "https://cdn.example.com/tplv-dy-water-v2/water-list.webp" |
| 855 | ], |
| 856 | }, |
| 857 | ] |
| 858 | }, |
| 859 | } |
| 860 | |
| 861 | urls = downloader._collect_image_urls(aweme_data) |
| 862 | |
| 863 | assert urls == [ |
| 864 | "https://cdn.example.com/clean-display.webp", |
| 865 | "https://cdn.example.com/clean-top.webp", |
| 866 | ] |
| 867 | |
| 868 | asyncio.run(api_client.close()) |
| 869 | |
| 870 | |
| 871 | def test_iter_gallery_items_image_list_key(tmp_path): |
| 872 | """Some responses use image_list instead of images.""" |
| 873 | downloader, api_client = _build_downloader(tmp_path) |
| 874 | |
| 875 | aweme_data = { |
| 876 | "aweme_id": "100004", |
| 877 | "image_post_info": { |
| 878 | "image_list": [{"display_image": {"url_list": ["https://example.com/img.webp"]}}] |
| 879 | }, |
| 880 | } |
| 881 | |
| 882 | items = downloader._iter_gallery_items(aweme_data) |
| 883 | assert len(items) == 1 |
| 884 | assert items[0]["display_image"]["url_list"][0] == "https://example.com/img.webp" |
| 885 | |
| 886 | asyncio.run(api_client.close()) |
| 887 | |
| 888 | |
| 889 | def test_iter_gallery_items_top_level_image_list(tmp_path): |
| 890 | """Fallback: top-level image_list key.""" |
| 891 | downloader, api_client = _build_downloader(tmp_path) |
| 892 | |
| 893 | aweme_data = { |
| 894 | "aweme_id": "100005", |
| 895 | "image_list": [{"url_list": ["https://example.com/top.webp"]}], |
| 896 | } |
| 897 | |
| 898 | items = downloader._iter_gallery_items(aweme_data) |
| 899 | assert len(items) == 1 |
| 900 | |
| 901 | asyncio.run(api_client.close()) |
| 902 |