| 1 | """Tests for task 2.2: downloader writes `author_sec_uid` (or None). |
| 2 | |
| 3 | Covers the spec requirement R12.12: |
| 4 | |
| 5 | > When a downloader persists an aweme row, the `author_sec_uid` column must |
| 6 | > carry `aweme.author.sec_uid` when present, and NULL otherwise. |
| 7 | |
| 8 | The tests are layered from narrow to broad: |
| 9 | |
| 10 | 1. ``extract_author_sec_uid`` is exercised directly on many payload shapes |
| 11 | (the defensive helper that the downloaders call at every `add_aweme` |
| 12 | site, see ``core/metadata.py``). |
| 13 | |
| 14 | 2. A ``VideoDownloader`` exercises the full ``_download_aweme_assets`` path |
| 15 | against a real on-disk SQLite ``Database`` (no mocks at the storage |
| 16 | boundary) with two aweme payloads — one with ``author.sec_uid`` set and |
| 17 | one without — and we assert the row in the ``aweme`` table matches. |
| 18 | |
| 19 | 3. A ``MusicDownloader`` test does the same for the music fallback path, |
| 20 | which is the other call site modified in task 2.1. |
| 21 | |
| 22 | The downloader tests avoid real network and disk media writes via the same |
| 23 | ``_download_with_retry`` / ``get_session`` monkeypatch pattern used across |
| 24 | ``tests/test_video_downloader.py`` and ``tests/test_music_downloader.py``. |
| 25 | """ |
| 26 | |
| 27 | from __future__ import annotations |
| 28 | |
| 29 | from typing import Any, Dict, Optional |
| 30 | |
| 31 | import pytest |
| 32 | |
| 33 | from auth import CookieManager |
| 34 | from config import ConfigLoader |
| 35 | from control import QueueManager, RateLimiter, RetryHandler |
| 36 | from core.api_client import DouyinAPIClient |
| 37 | from core.metadata import extract_author_sec_uid |
| 38 | from core.music_downloader import MusicDownloader |
| 39 | from core.video_downloader import VideoDownloader |
| 40 | from storage import Database, FileManager |
| 41 | |
| 42 | |
| 43 | # --------------------------------------------------------------------------- |
| 44 | # 1. Pure helper — extract_author_sec_uid |
| 45 | # --------------------------------------------------------------------------- |
| 46 | def test_extract_returns_sec_uid_when_present(): |
| 47 | assert extract_author_sec_uid({"author": {"sec_uid": "SEC_X"}}) == "SEC_X" |
| 48 | |
| 49 | |
| 50 | def test_extract_strips_whitespace(): |
| 51 | assert extract_author_sec_uid({"author": {"sec_uid": " SEC_Y "}}) == "SEC_Y" |
| 52 | |
| 53 | |
| 54 | @pytest.mark.parametrize( |
| 55 | "payload", |
| 56 | [ |
| 57 | pytest.param(None, id="input-is-none"), |
| 58 | pytest.param("not-a-mapping", id="input-is-string"), |
| 59 | pytest.param({}, id="author-missing"), |
| 60 | pytest.param({"author": None}, id="author-is-none"), |
| 61 | pytest.param({"author": "not-a-mapping"}, id="author-is-string"), |
| 62 | pytest.param({"author": {}}, id="sec-uid-missing"), |
| 63 | pytest.param({"author": {"sec_uid": None}}, id="sec-uid-is-none"), |
| 64 | pytest.param({"author": {"sec_uid": 123}}, id="sec-uid-not-string"), |
| 65 | pytest.param({"author": {"sec_uid": ""}}, id="sec-uid-empty-string"), |
| 66 | pytest.param({"author": {"sec_uid": " "}}, id="sec-uid-whitespace"), |
| 67 | ], |
| 68 | ) |
| 69 | def test_extract_returns_none_for_invalid_payloads(payload): |
| 70 | assert extract_author_sec_uid(payload) is None |
| 71 | |
| 72 | |
| 73 | # --------------------------------------------------------------------------- |
| 74 | # 2. VideoDownloader end-to-end: the row in `aweme` carries author_sec_uid |
| 75 | # --------------------------------------------------------------------------- |
| 76 | def _build_video_downloader( |
| 77 | tmp_path, database: Database |
| 78 | ) -> tuple[VideoDownloader, DouyinAPIClient]: |
| 79 | config = ConfigLoader() |
| 80 | # Disable every optional side-car asset so the test only drives the |
| 81 | # video + db.add_aweme path. |
| 82 | config.update( |
| 83 | path=str(tmp_path), |
| 84 | music=False, |
| 85 | cover=False, |
| 86 | avatar=False, |
| 87 | json=False, |
| 88 | folderstyle=True, |
| 89 | transcript={"enabled": False}, |
| 90 | ) |
| 91 | file_manager = FileManager(str(tmp_path)) |
| 92 | cookie_manager = CookieManager(str(tmp_path / ".cookies.json")) |
| 93 | api_client = DouyinAPIClient({}) |
| 94 | |
| 95 | downloader = VideoDownloader( |
| 96 | config, |
| 97 | api_client, |
| 98 | file_manager, |
| 99 | cookie_manager, |
| 100 | database=database, |
| 101 | rate_limiter=RateLimiter(max_per_second=10), |
| 102 | retry_handler=RetryHandler(max_retries=1), |
| 103 | queue_manager=QueueManager(max_workers=1), |
| 104 | ) |
| 105 | return downloader, api_client |
| 106 | |
| 107 | |
| 108 | async def _fetch_db_row(db: Database, aweme_id: str) -> Optional[Dict[str, Any]]: |
| 109 | conn = await db._get_conn() |
| 110 | cursor = await conn.execute( |
| 111 | "SELECT aweme_id, author_id, author_sec_uid FROM aweme WHERE aweme_id = ?", |
| 112 | (aweme_id,), |
| 113 | ) |
| 114 | row = await cursor.fetchone() |
| 115 | if row is None: |
| 116 | return None |
| 117 | return {"aweme_id": row[0], "author_id": row[1], "author_sec_uid": row[2]} |
| 118 | |
| 119 | |
| 120 | async def test_video_downloader_persists_author_sec_uid_when_present(tmp_path, monkeypatch): |
| 121 | db = Database(db_path=str(tmp_path / "test.db")) |
| 122 | await db.initialize() |
| 123 | try: |
| 124 | downloader, api_client = _build_video_downloader(tmp_path, db) |
| 125 | |
| 126 | async def _fake_get_session(): |
| 127 | return object() |
| 128 | |
| 129 | monkeypatch.setattr(api_client, "get_session", _fake_get_session) |
| 130 | |
| 131 | async def _fake_download_with_retry(self, _url, _save_path, _session, **_kwargs): |
| 132 | return True |
| 133 | |
| 134 | downloader._download_with_retry = _fake_download_with_retry.__get__( |
| 135 | downloader, VideoDownloader |
| 136 | ) |
| 137 | |
| 138 | aweme_id = "7600224486650121526" |
| 139 | aweme_data = { |
| 140 | "aweme_id": aweme_id, |
| 141 | "desc": "has sec_uid", |
| 142 | "create_time": 1707303025, |
| 143 | "author": { |
| 144 | "uid": "u1", |
| 145 | "nickname": "Alice", |
| 146 | "sec_uid": "SEC_X", |
| 147 | }, |
| 148 | "video": {"play_addr": {"url_list": ["https://example.com/video.mp4"]}}, |
| 149 | } |
| 150 | |
| 151 | success = await downloader._download_aweme_assets( |
| 152 | aweme_data, author_name="Alice", mode="post" |
| 153 | ) |
| 154 | assert success is True |
| 155 | |
| 156 | row = await _fetch_db_row(db, aweme_id) |
| 157 | assert row is not None, "expected aweme row to be persisted" |
| 158 | assert row["author_sec_uid"] == "SEC_X" |
| 159 | # sanity-check that the row is otherwise well-formed |
| 160 | assert row["author_id"] == "u1" |
| 161 | |
| 162 | await api_client.close() |
| 163 | finally: |
| 164 | await db.close() |
| 165 | |
| 166 | |
| 167 | async def test_video_downloader_persists_null_when_sec_uid_missing(tmp_path, monkeypatch): |
| 168 | db = Database(db_path=str(tmp_path / "test.db")) |
| 169 | await db.initialize() |
| 170 | try: |
| 171 | downloader, api_client = _build_video_downloader(tmp_path, db) |
| 172 | |
| 173 | async def _fake_get_session(): |
| 174 | return object() |
| 175 | |
| 176 | monkeypatch.setattr(api_client, "get_session", _fake_get_session) |
| 177 | |
| 178 | async def _fake_download_with_retry(self, _url, _save_path, _session, **_kwargs): |
| 179 | return True |
| 180 | |
| 181 | downloader._download_with_retry = _fake_download_with_retry.__get__( |
| 182 | downloader, VideoDownloader |
| 183 | ) |
| 184 | |
| 185 | aweme_id = "7600224486650121999" |
| 186 | aweme_data = { |
| 187 | "aweme_id": aweme_id, |
| 188 | "desc": "missing sec_uid", |
| 189 | "create_time": 1707303025, |
| 190 | "author": {"uid": "u2", "nickname": "Bob"}, # no sec_uid |
| 191 | "video": {"play_addr": {"url_list": ["https://example.com/video.mp4"]}}, |
| 192 | } |
| 193 | |
| 194 | success = await downloader._download_aweme_assets( |
| 195 | aweme_data, author_name="Bob", mode="post" |
| 196 | ) |
| 197 | assert success is True |
| 198 | |
| 199 | row = await _fetch_db_row(db, aweme_id) |
| 200 | assert row is not None |
| 201 | assert row["author_sec_uid"] is None |
| 202 | assert row["author_id"] == "u2" |
| 203 | |
| 204 | await api_client.close() |
| 205 | finally: |
| 206 | await db.close() |
| 207 | |
| 208 | |
| 209 | async def test_video_downloader_persists_null_when_author_absent(tmp_path, monkeypatch): |
| 210 | """Entirely missing `author` object ⇒ NULL (not an exception).""" |
| 211 | db = Database(db_path=str(tmp_path / "test.db")) |
| 212 | await db.initialize() |
| 213 | try: |
| 214 | downloader, api_client = _build_video_downloader(tmp_path, db) |
| 215 | |
| 216 | async def _fake_get_session(): |
| 217 | return object() |
| 218 | |
| 219 | monkeypatch.setattr(api_client, "get_session", _fake_get_session) |
| 220 | |
| 221 | async def _fake_download_with_retry(self, _url, _save_path, _session, **_kwargs): |
| 222 | return True |
| 223 | |
| 224 | downloader._download_with_retry = _fake_download_with_retry.__get__( |
| 225 | downloader, VideoDownloader |
| 226 | ) |
| 227 | |
| 228 | aweme_id = "7600224486650122020" |
| 229 | aweme_data = { |
| 230 | "aweme_id": aweme_id, |
| 231 | "desc": "no author object", |
| 232 | "create_time": 1707303025, |
| 233 | "video": {"play_addr": {"url_list": ["https://example.com/video.mp4"]}}, |
| 234 | } |
| 235 | |
| 236 | success = await downloader._download_aweme_assets( |
| 237 | aweme_data, author_name="anon", mode="post" |
| 238 | ) |
| 239 | assert success is True |
| 240 | |
| 241 | row = await _fetch_db_row(db, aweme_id) |
| 242 | assert row is not None |
| 243 | assert row["author_sec_uid"] is None |
| 244 | |
| 245 | await api_client.close() |
| 246 | finally: |
| 247 | await db.close() |
| 248 | |
| 249 | |
| 250 | # --------------------------------------------------------------------------- |
| 251 | # 3. MusicDownloader: covers the other call site touched by task 2.1 |
| 252 | # --------------------------------------------------------------------------- |
| 253 | class _MusicAPIClient: |
| 254 | BASE_URL = "https://www.douyin.com" |
| 255 | headers = {"User-Agent": "UnitTestAgent/1.0"} |
| 256 | |
| 257 | def __init__(self, detail: Dict[str, Any]): |
| 258 | self._detail = detail |
| 259 | |
| 260 | async def get_music_detail(self, _music_id: str): |
| 261 | return self._detail |
| 262 | |
| 263 | async def get_session(self): |
| 264 | return object() |
| 265 | |
| 266 | |
| 267 | def _build_music_downloader( |
| 268 | tmp_path, database: Database, api_client: _MusicAPIClient |
| 269 | ) -> MusicDownloader: |
| 270 | config = ConfigLoader() |
| 271 | config.update(path=str(tmp_path), cover=False, json=False) |
| 272 | file_manager = FileManager(str(tmp_path)) |
| 273 | return MusicDownloader( |
| 274 | config=config, |
| 275 | api_client=api_client, |
| 276 | file_manager=file_manager, |
| 277 | cookie_manager=CookieManager(str(tmp_path / ".cookies.json")), |
| 278 | database=database, |
| 279 | rate_limiter=RateLimiter(max_per_second=10), |
| 280 | retry_handler=RetryHandler(max_retries=1), |
| 281 | queue_manager=QueueManager(max_workers=1), |
| 282 | ) |
| 283 | |
| 284 | |
| 285 | async def test_music_downloader_persists_author_sec_uid_when_present(tmp_path, monkeypatch): |
| 286 | db = Database(db_path=str(tmp_path / "test.db")) |
| 287 | await db.initialize() |
| 288 | try: |
| 289 | detail = { |
| 290 | "title": "song-a", |
| 291 | "author_name": "artist-a", |
| 292 | "author": {"sec_uid": "SEC_MUSIC"}, |
| 293 | "play_url": {"url_list": ["https://example.com/music.mp3"]}, |
| 294 | } |
| 295 | api_client = _MusicAPIClient(detail) |
| 296 | downloader = _build_music_downloader(tmp_path, db, api_client) |
| 297 | |
| 298 | async def _fake_download_with_retry(self, _url, _save_path, _session, **_kwargs): |
| 299 | return True |
| 300 | |
| 301 | monkeypatch.setattr( |
| 302 | downloader, |
| 303 | "_download_with_retry", |
| 304 | _fake_download_with_retry.__get__(downloader, MusicDownloader), |
| 305 | ) |
| 306 | |
| 307 | result = await downloader.download({"music_id": "7600"}) |
| 308 | assert result.success == 1 |
| 309 | |
| 310 | row = await _fetch_db_row(db, "music_7600") |
| 311 | assert row is not None |
| 312 | assert row["author_sec_uid"] == "SEC_MUSIC" |
| 313 | finally: |
| 314 | await db.close() |
| 315 | |
| 316 | |
| 317 | async def test_music_downloader_persists_null_when_sec_uid_missing(tmp_path, monkeypatch): |
| 318 | db = Database(db_path=str(tmp_path / "test.db")) |
| 319 | await db.initialize() |
| 320 | try: |
| 321 | detail = { |
| 322 | "title": "song-b", |
| 323 | "author_name": "artist-b", |
| 324 | # no author.sec_uid |
| 325 | "play_url": {"url_list": ["https://example.com/music.mp3"]}, |
| 326 | } |
| 327 | api_client = _MusicAPIClient(detail) |
| 328 | downloader = _build_music_downloader(tmp_path, db, api_client) |
| 329 | |
| 330 | async def _fake_download_with_retry(self, _url, _save_path, _session, **_kwargs): |
| 331 | return True |
| 332 | |
| 333 | monkeypatch.setattr( |
| 334 | downloader, |
| 335 | "_download_with_retry", |
| 336 | _fake_download_with_retry.__get__(downloader, MusicDownloader), |
| 337 | ) |
| 338 | |
| 339 | result = await downloader.download({"music_id": "7601"}) |
| 340 | assert result.success == 1 |
| 341 | |
| 342 | row = await _fetch_db_row(db, "music_7601") |
| 343 | assert row is not None |
| 344 | assert row["author_sec_uid"] is None |
| 345 | finally: |
| 346 | await db.close() |
| 347 |