| 1 | import pytest |
| 2 | |
| 3 | from auth import CookieManager |
| 4 | from config import ConfigLoader |
| 5 | from control import QueueManager, RateLimiter, RetryHandler |
| 6 | from core.music_downloader import MusicDownloader |
| 7 | from storage import FileManager |
| 8 | |
| 9 | |
| 10 | class _FakeAPIClient: |
| 11 | BASE_URL = "https://www.douyin.com" |
| 12 | headers = {"User-Agent": "UnitTestAgent/1.0"} |
| 13 | |
| 14 | async def get_music_detail(self, _music_id: str): |
| 15 | return { |
| 16 | "title": "test-music", |
| 17 | "author_name": "test-author", |
| 18 | "play_url": {"url_list": ["https://example.com/music.mp3"]}, |
| 19 | } |
| 20 | |
| 21 | async def get_session(self): |
| 22 | return object() |
| 23 | |
| 24 | |
| 25 | @pytest.mark.asyncio |
| 26 | async def test_music_downloader_downloads_music_asset(tmp_path, monkeypatch): |
| 27 | config = ConfigLoader() |
| 28 | config.update(path=str(tmp_path), cover=False, json=False) |
| 29 | file_manager = FileManager(str(tmp_path)) |
| 30 | downloader = MusicDownloader( |
| 31 | config=config, |
| 32 | api_client=_FakeAPIClient(), |
| 33 | file_manager=file_manager, |
| 34 | cookie_manager=CookieManager(str(tmp_path / ".cookies.json")), |
| 35 | database=None, |
| 36 | rate_limiter=RateLimiter(max_per_second=10), |
| 37 | retry_handler=RetryHandler(max_retries=1), |
| 38 | queue_manager=QueueManager(max_workers=1), |
| 39 | ) |
| 40 | |
| 41 | saved_paths = [] |
| 42 | |
| 43 | async def _fake_download_with_retry(self, _url, save_path, _session, **_kwargs): |
| 44 | saved_paths.append(save_path) |
| 45 | return True |
| 46 | |
| 47 | monkeypatch.setattr( |
| 48 | downloader, |
| 49 | "_download_with_retry", |
| 50 | _fake_download_with_retry.__get__(downloader, MusicDownloader), |
| 51 | ) |
| 52 | |
| 53 | result = await downloader.download({"music_id": "7600224486650121999"}) |
| 54 | |
| 55 | assert result.total == 1 |
| 56 | assert result.success == 1 |
| 57 | assert any(path.suffix == ".mp3" for path in saved_paths) |
| 58 | |
| 59 | |
| 60 | @pytest.mark.asyncio |
| 61 | async def test_music_downloader_uses_extension_from_music_url(tmp_path, monkeypatch): |
| 62 | class _FakeM4AAPIClient(_FakeAPIClient): |
| 63 | async def get_music_detail(self, _music_id: str): |
| 64 | return { |
| 65 | "title": "test-music", |
| 66 | "author_name": "test-author", |
| 67 | "play_url": {"url_list": ["https://example.com/music_track.m4a?x=1"]}, |
| 68 | } |
| 69 | |
| 70 | config = ConfigLoader() |
| 71 | config.update(path=str(tmp_path), cover=False, json=False) |
| 72 | file_manager = FileManager(str(tmp_path)) |
| 73 | downloader = MusicDownloader( |
| 74 | config=config, |
| 75 | api_client=_FakeM4AAPIClient(), |
| 76 | file_manager=file_manager, |
| 77 | cookie_manager=CookieManager(str(tmp_path / ".cookies.json")), |
| 78 | database=None, |
| 79 | rate_limiter=RateLimiter(max_per_second=10), |
| 80 | retry_handler=RetryHandler(max_retries=1), |
| 81 | queue_manager=QueueManager(max_workers=1), |
| 82 | ) |
| 83 | |
| 84 | saved_paths = [] |
| 85 | |
| 86 | async def _fake_download_with_retry(self, _url, save_path, _session, **_kwargs): |
| 87 | saved_paths.append(save_path) |
| 88 | return True |
| 89 | |
| 90 | monkeypatch.setattr( |
| 91 | downloader, |
| 92 | "_download_with_retry", |
| 93 | _fake_download_with_retry.__get__(downloader, MusicDownloader), |
| 94 | ) |
| 95 | |
| 96 | result = await downloader.download({"music_id": "7600224486650122000"}) |
| 97 | |
| 98 | assert result.success == 1 |
| 99 | assert any(path.suffix == ".m4a" for path in saved_paths) |
| 100 | |
| 101 | |
| 102 | @pytest.mark.asyncio |
| 103 | async def test_music_downloader_falls_back_to_first_aweme_when_direct_audio_missing( |
| 104 | tmp_path, monkeypatch |
| 105 | ): |
| 106 | class _FallbackAPIClient(_FakeAPIClient): |
| 107 | async def get_music_detail(self, _music_id: str): |
| 108 | return { |
| 109 | "title": "fallback-music", |
| 110 | "author_name": "fallback-author", |
| 111 | } |
| 112 | |
| 113 | async def get_music_aweme(self, _music_id: str, cursor: int = 0, count: int = 1): |
| 114 | assert cursor == 0 |
| 115 | assert count == 1 |
| 116 | return { |
| 117 | "items": [ |
| 118 | { |
| 119 | "aweme_id": "fallback-aweme-1", |
| 120 | "author": {"nickname": "fallback-author"}, |
| 121 | "video": {"play_addr": {"url_list": ["https://example.com/video.mp4"]}}, |
| 122 | } |
| 123 | ] |
| 124 | } |
| 125 | |
| 126 | config = ConfigLoader() |
| 127 | config.update(path=str(tmp_path), cover=False, json=False) |
| 128 | file_manager = FileManager(str(tmp_path)) |
| 129 | downloader = MusicDownloader( |
| 130 | config=config, |
| 131 | api_client=_FallbackAPIClient(), |
| 132 | file_manager=file_manager, |
| 133 | cookie_manager=CookieManager(str(tmp_path / ".cookies.json")), |
| 134 | database=None, |
| 135 | rate_limiter=RateLimiter(max_per_second=10), |
| 136 | retry_handler=RetryHandler(max_retries=1), |
| 137 | queue_manager=QueueManager(max_workers=1), |
| 138 | ) |
| 139 | |
| 140 | downloaded_awemes = [] |
| 141 | |
| 142 | async def _fake_should_download(self, _aweme_id: str): |
| 143 | return True |
| 144 | |
| 145 | async def _fake_download_aweme_assets(self, aweme_data, author_name, mode=None): |
| 146 | downloaded_awemes.append((aweme_data["aweme_id"], author_name, mode)) |
| 147 | return True |
| 148 | |
| 149 | monkeypatch.setattr( |
| 150 | downloader, |
| 151 | "_should_download", |
| 152 | _fake_should_download.__get__(downloader, MusicDownloader), |
| 153 | ) |
| 154 | monkeypatch.setattr( |
| 155 | downloader, |
| 156 | "_download_aweme_assets", |
| 157 | _fake_download_aweme_assets.__get__(downloader, MusicDownloader), |
| 158 | ) |
| 159 | |
| 160 | result = await downloader.download({"music_id": "7600224486650122001"}) |
| 161 | |
| 162 | assert result.total == 1 |
| 163 | assert result.success == 1 |
| 164 | assert downloaded_awemes == [("fallback-aweme-1", "fallback-author", "music")] |
| 165 |