返回 douyin-downloader
test_user_mode_strategies.py
根目录 / tests / test_user_mode_strategies.py
1 import asyncio
2
3 from core.user_modes.collect_mix_strategy import CollectMixUserModeStrategy
4 from core.user_modes.collect_strategy import CollectUserModeStrategy
5 from core.user_modes.like_strategy import LikeUserModeStrategy
6 from core.user_modes.mix_strategy import MixUserModeStrategy
7 from core.user_modes.music_strategy import MusicUserModeStrategy
8 from core.user_modes.post_strategy import PostUserModeStrategy
9
10
11 class _NoopRateLimiter:
12 async def acquire(self):
13 return
14
15
16 def _make_aweme(aweme_id: str):
17 return {
18 "aweme_id": aweme_id,
19 "create_time": 1700000000,
20 "video": {"play_addr": {"url_list": ["https://example.com/video.mp4"]}},
21 }
22
23
24 def test_like_strategy_collects_items_from_api():
25 class _API:
26 async def get_user_like(self, _sec_uid, max_cursor=0, count=20):
27 if max_cursor > 0:
28 return {"items": [], "has_more": False, "max_cursor": max_cursor}
29 return {"items": [_make_aweme("111")], "has_more": False, "max_cursor": 0}
30
31 class _Downloader:
32 def __init__(self):
33 self.api_client = _API()
34 self.rate_limiter = _NoopRateLimiter()
35 self.config = type(
36 "Cfg",
37 (),
38 {
39 "get": lambda _self, key, default=None: {
40 "number": {"like": 0},
41 "increase": {"like": False},
42 }.get(key, default)
43 },
44 )()
45 self.database = None
46 self._filter_by_time = lambda items: items
47 self._limit_count = lambda items, _mode: items
48
49 strategy = LikeUserModeStrategy(_Downloader())
50 items = asyncio.run(strategy.collect_items("sec_uid_x", {"uid": "uid-1"}))
51 assert [item["aweme_id"] for item in items] == ["111"]
52
53
54 def test_like_strategy_increment_stops_at_first_downloaded_aweme():
55 class _API:
56 def __init__(self):
57 self.calls = []
58
59 async def get_user_like(self, _sec_uid, max_cursor=0, count=20):
60 self.calls.append(max_cursor)
61 if max_cursor == 0:
62 return {
63 "items": [_make_aweme("new-1"), _make_aweme("old-1")],
64 "has_more": True,
65 "max_cursor": 1,
66 }
67 return {
68 "items": [_make_aweme("older-1")],
69 "has_more": False,
70 "max_cursor": max_cursor,
71 }
72
73 class _Database:
74 async def get_latest_aweme_time(self, _author_id):
75 return None
76
77 async def is_downloaded(self, aweme_id):
78 return aweme_id == "old-1"
79
80 class _Downloader:
81 def __init__(self):
82 self.api_client = _API()
83 self.rate_limiter = _NoopRateLimiter()
84 self.config = type(
85 "Cfg",
86 (),
87 {
88 "get": lambda _self, key, default=None: {
89 "number": {"like": 0},
90 "increase": {"like": True},
91 }.get(key, default)
92 },
93 )()
94 self.database = _Database()
95 self._filter_by_time = lambda items: items
96 self._limit_count = lambda items, _mode: items
97
98 downloader = _Downloader()
99 strategy = LikeUserModeStrategy(downloader)
100 items = asyncio.run(strategy.collect_items("sec_uid_x", {"uid": "uid-1"}))
101
102 assert [item["aweme_id"] for item in items] == ["new-1"]
103 assert downloader.api_client.calls == [0]
104
105
106 def test_post_strategy_calls_browser_recover_when_pagination_restricted():
107 class _API:
108 async def get_user_post(self, _sec_uid, max_cursor=0, count=20):
109 if max_cursor == 0:
110 return {
111 "items": [_make_aweme("111")],
112 "has_more": True,
113 "max_cursor": 123,
114 "status_code": 0,
115 }
116 return {"items": [], "has_more": False, "max_cursor": max_cursor, "status_code": 0}
117
118 class _Downloader:
119 def __init__(self):
120 self.api_client = _API()
121 self.rate_limiter = _NoopRateLimiter()
122 self.database = None
123 self.config = type(
124 "Cfg",
125 (),
126 {
127 "get": lambda _self, key, default=None: {
128 "number": {"post": 0},
129 "increase": {"post": False},
130 "browser_fallback": {"enabled": True},
131 }.get(key, default)
132 },
133 )()
134 self.recovered_called = False
135 self._progress_update_step = lambda *_args, **_kwargs: None
136 self._filter_by_time = lambda items: items
137 self._limit_count = lambda items, _mode: items
138
139 async def _recover_user_post_with_browser(self, sec_uid, user_info, aweme_list):
140 self.recovered_called = True
141 aweme_list.append(_make_aweme("222"))
142
143 downloader = _Downloader()
144 strategy = PostUserModeStrategy(downloader)
145 items = asyncio.run(strategy.collect_items("sec_uid_x", {"uid": "uid-1"}))
146
147 assert downloader.recovered_called is True
148 assert [item["aweme_id"] for item in items] == ["111", "222"]
149
150
151 def test_post_strategy_calls_browser_recover_when_cursor_stalls():
152 class _API:
153 async def get_user_post(self, _sec_uid, max_cursor=0, count=20):
154 return {
155 "items": [_make_aweme("333")],
156 "has_more": True,
157 "max_cursor": max_cursor,
158 "status_code": 0,
159 }
160
161 class _Downloader:
162 def __init__(self):
163 self.api_client = _API()
164 self.rate_limiter = _NoopRateLimiter()
165 self.database = None
166 self.config = type(
167 "Cfg",
168 (),
169 {
170 "get": lambda _self, key, default=None: {
171 "number": {"post": 0},
172 "increase": {"post": False},
173 "browser_fallback": {"enabled": True},
174 }.get(key, default)
175 },
176 )()
177 self.recovered_called = False
178 self._progress_update_step = lambda *_args, **_kwargs: None
179 self._filter_by_time = lambda items: items
180 self._limit_count = lambda items, _mode: items
181
182 async def _recover_user_post_with_browser(self, sec_uid, user_info, aweme_list):
183 self.recovered_called = True
184 aweme_list.append(_make_aweme("444"))
185
186 downloader = _Downloader()
187 strategy = PostUserModeStrategy(downloader)
188 items = asyncio.run(strategy.collect_items("sec_uid_x", {"uid": "uid-1"}))
189
190 assert downloader.recovered_called is True
191 assert [item["aweme_id"] for item in items] == ["333", "444"]
192
193
194 def test_mix_strategy_filters_partial_aweme_items_without_metadata_inflation():
195 class _API:
196 async def get_user_mix(self, _sec_uid, max_cursor=0, count=20):
197 return {
198 "items": [
199 {"aweme_id": "111"},
200 {"mix_info": {"mix_id": "mix-only-meta"}},
201 ],
202 "has_more": False,
203 "max_cursor": 0,
204 }
205
206 class _Downloader:
207 def __init__(self):
208 self.api_client = _API()
209 self.rate_limiter = _NoopRateLimiter()
210 self.database = None
211 self.config = type(
212 "Cfg",
213 (),
214 {
215 "get": lambda _self, key, default=None: {
216 "number": {"mix": 0},
217 "increase": {"mix": False},
218 }.get(key, default)
219 },
220 )()
221 self._filter_by_time = lambda items: items
222 self._limit_count = lambda items, _mode: items
223
224 strategy = MixUserModeStrategy(_Downloader())
225 items = asyncio.run(strategy.collect_items("sec_uid_x", {"uid": "uid-1"}))
226 assert items == [{"aweme_id": "111"}]
227
228
229 def test_mix_strategy_expansion_does_not_apply_number_limit_early():
230 class _API:
231 async def get_user_mix(self, _sec_uid, max_cursor=0, count=20):
232 return {
233 "items": [{"mix_info": {"mix_id": "mix-1"}}],
234 "has_more": False,
235 "max_cursor": 0,
236 }
237
238 async def get_mix_aweme(self, _mix_id, cursor=0, count=20):
239 return {
240 "items": [{"aweme_id": "m-1"}, {"aweme_id": "m-2"}],
241 "has_more": False,
242 "max_cursor": 0,
243 }
244
245 class _Downloader:
246 def __init__(self):
247 self.api_client = _API()
248 self.rate_limiter = _NoopRateLimiter()
249 self.database = None
250 self.config = type(
251 "Cfg",
252 (),
253 {
254 "get": lambda _self, key, default=None: {
255 "number": {"mix": 1},
256 "increase": {"mix": False},
257 }.get(key, default)
258 },
259 )()
260 self._filter_by_time = lambda items: items
261 self._limit_count = lambda items, _mode: items
262
263 strategy = MixUserModeStrategy(_Downloader())
264 items = asyncio.run(strategy.collect_items("sec_uid_x", {"uid": "uid-1"}))
265 assert [item["aweme_id"] for item in items] == ["m-1", "m-2"]
266
267
268 def test_music_strategy_filters_partial_aweme_items_without_metadata_inflation():
269 class _API:
270 async def get_user_music(self, _sec_uid, max_cursor=0, count=20):
271 return {
272 "items": [
273 {"aweme_id": "222"},
274 {"music_info": {"id": "music-only-meta"}},
275 ],
276 "has_more": False,
277 "max_cursor": 0,
278 }
279
280 class _Downloader:
281 def __init__(self):
282 self.api_client = _API()
283 self.rate_limiter = _NoopRateLimiter()
284 self.database = None
285 self.config = type(
286 "Cfg",
287 (),
288 {
289 "get": lambda _self, key, default=None: {
290 "number": {"music": 0},
291 "increase": {"music": False},
292 }.get(key, default)
293 },
294 )()
295 self._filter_by_time = lambda items: items
296 self._limit_count = lambda items, _mode: items
297
298 strategy = MusicUserModeStrategy(_Downloader())
299 items = asyncio.run(strategy.collect_items("sec_uid_x", {"uid": "uid-1"}))
300 assert items == [{"aweme_id": "222"}]
301
302
303 def test_music_strategy_expansion_does_not_apply_number_limit_early():
304 class _API:
305 async def get_user_music(self, _sec_uid, max_cursor=0, count=20):
306 return {
307 "items": [{"music_info": {"id": "music-1"}}],
308 "has_more": False,
309 "max_cursor": 0,
310 }
311
312 async def get_music_aweme(self, _music_id, cursor=0, count=20):
313 return {
314 "items": [{"aweme_id": "mu-1"}, {"aweme_id": "mu-2"}],
315 "has_more": False,
316 "max_cursor": 0,
317 }
318
319 class _Downloader:
320 def __init__(self):
321 self.api_client = _API()
322 self.rate_limiter = _NoopRateLimiter()
323 self.database = None
324 self.config = type(
325 "Cfg",
326 (),
327 {
328 "get": lambda _self, key, default=None: {
329 "number": {"music": 1},
330 "increase": {"music": False},
331 }.get(key, default)
332 },
333 )()
334 self._filter_by_time = lambda items: items
335 self._limit_count = lambda items, _mode: items
336
337 strategy = MusicUserModeStrategy(_Downloader())
338 items = asyncio.run(strategy.collect_items("sec_uid_x", {"uid": "uid-1"}))
339 assert [item["aweme_id"] for item in items] == ["mu-1", "mu-2"]
340
341
342 def test_collect_strategy_expands_collect_folders_and_deduplicates_aweme():
343 class _API:
344 async def get_user_collects(self, _sec_uid, max_cursor=0, count=20):
345 if max_cursor > 0:
346 return {"items": [], "has_more": False, "max_cursor": max_cursor}
347 return {
348 "items": [
349 {"collects_id_str": "collect-1"},
350 {"collects_id_str": "collect-2"},
351 ],
352 "has_more": False,
353 "max_cursor": 0,
354 }
355
356 async def get_collect_aweme(self, collects_id, max_cursor=0, count=20):
357 assert max_cursor == 0
358 if collects_id == "collect-1":
359 return {
360 "items": [{"aweme_id": "c-1"}, {"aweme_id": "dup"}],
361 "has_more": False,
362 "max_cursor": 0,
363 }
364 return {
365 "items": [{"aweme_id": "dup"}, {"aweme_id": "c-2"}],
366 "has_more": False,
367 "max_cursor": 0,
368 }
369
370 class _Downloader:
371 def __init__(self):
372 self.api_client = _API()
373 self.rate_limiter = _NoopRateLimiter()
374 self.database = None
375 self.config = type(
376 "Cfg",
377 (),
378 {
379 "get": lambda _self, key, default=None: {
380 "number": {"collect": 0},
381 "increase": {"collect": False},
382 }.get(key, default)
383 },
384 )()
385 self._filter_by_time = lambda items: items
386 self._limit_count = lambda items, _mode: items
387
388 strategy = CollectUserModeStrategy(_Downloader())
389 items = asyncio.run(strategy.collect_items("self", {"uid": "self"}))
390 assert [item["aweme_id"] for item in items] == ["c-1", "dup", "c-2"]
391
392
393 def test_collect_strategy_expansion_does_not_apply_number_limit_or_increase_early():
394 class _Database:
395 async def get_latest_aweme_time(self, _author_id):
396 return 1700000000
397
398 class _API:
399 async def get_user_collects(self, _sec_uid, max_cursor=0, count=20):
400 return {
401 "items": [
402 {"collects_id_str": "collect-1"},
403 {"collects_id_str": "collect-2"},
404 ],
405 "has_more": False,
406 "max_cursor": 0,
407 }
408
409 async def get_collect_aweme(self, collects_id, max_cursor=0, count=20):
410 if collects_id == "collect-1":
411 return {
412 "items": [{"aweme_id": "c-1", "create_time": 1700000001}],
413 "has_more": False,
414 "max_cursor": 0,
415 }
416 return {
417 "items": [{"aweme_id": "c-2", "create_time": 1700000002}],
418 "has_more": False,
419 "max_cursor": 0,
420 }
421
422 class _Downloader:
423 def __init__(self):
424 self.api_client = _API()
425 self.rate_limiter = _NoopRateLimiter()
426 self.database = _Database()
427 self.config = type(
428 "Cfg",
429 (),
430 {
431 "get": lambda _self, key, default=None: {
432 "number": {"collect": 1},
433 "increase": {"collect": True},
434 }.get(key, default)
435 },
436 )()
437 self._filter_by_time = lambda items: items
438 self._limit_count = lambda items, _mode: items
439
440 strategy = CollectUserModeStrategy(_Downloader())
441 items = asyncio.run(strategy.collect_items("self", {"uid": "self"}))
442 assert [item["aweme_id"] for item in items] == ["c-1", "c-2"]
443
444
445 def test_collect_mix_strategy_expands_collected_mix_items():
446 class _API:
447 async def get_user_collect_mix(self, _sec_uid, max_cursor=0, count=20):
448 return {
449 "items": [
450 {"mix_info": {"mix_id": "mix-1"}},
451 {"mix_id": "mix-2"},
452 ],
453 "has_more": False,
454 "max_cursor": 0,
455 }
456
457 async def get_mix_aweme(self, mix_id, cursor=0, count=20):
458 if mix_id == "mix-1":
459 return {
460 "items": [{"aweme_id": "mix-aweme-1"}],
461 "has_more": False,
462 "max_cursor": 0,
463 }
464 return {
465 "items": [{"aweme_id": "mix-aweme-2"}],
466 "has_more": False,
467 "max_cursor": 0,
468 }
469
470 class _Downloader:
471 def __init__(self):
472 self.api_client = _API()
473 self.rate_limiter = _NoopRateLimiter()
474 self.database = None
475 self.config = type(
476 "Cfg",
477 (),
478 {
479 "get": lambda _self, key, default=None: {
480 "number": {"collectmix": 0},
481 "increase": {"collectmix": False},
482 }.get(key, default)
483 },
484 )()
485 self._filter_by_time = lambda items: items
486 self._limit_count = lambda items, _mode: items
487
488 strategy = CollectMixUserModeStrategy(_Downloader())
489 items = asyncio.run(strategy.collect_items("self", {"uid": "self"}))
490 assert [item["aweme_id"] for item in items] == ["mix-aweme-1", "mix-aweme-2"]
491
492
493 def test_collect_mix_strategy_keeps_direct_aweme_items_and_expands_remaining_metadata():
494 class _API:
495 async def get_user_collect_mix(self, _sec_uid, max_cursor=0, count=20):
496 return {
497 "items": [
498 {"aweme_id": "mix-preview-1"},
499 {"mix_info": {"mix_id": "mix-1"}},
500 ],
501 "has_more": False,
502 "max_cursor": 0,
503 }
504
505 async def get_mix_aweme(self, mix_id, cursor=0, count=20):
506 assert mix_id == "mix-1"
507 return {
508 "items": [{"aweme_id": "mix-aweme-1"}],
509 "has_more": False,
510 "max_cursor": 0,
511 }
512
513 class _Downloader:
514 def __init__(self):
515 self.api_client = _API()
516 self.rate_limiter = _NoopRateLimiter()
517 self.database = None
518 self.config = type(
519 "Cfg",
520 (),
521 {
522 "get": lambda _self, key, default=None: {
523 "number": {"collectmix": 0},
524 "increase": {"collectmix": False},
525 }.get(key, default)
526 },
527 )()
528 self._filter_by_time = lambda items: items
529 self._limit_count = lambda items, _mode: items
530
531 strategy = CollectMixUserModeStrategy(_Downloader())
532 items = asyncio.run(strategy.collect_items("self", {"uid": "self"}))
533 assert [item["aweme_id"] for item in items] == ["mix-preview-1", "mix-aweme-1"]
534
535
536 def test_collect_mix_strategy_expansion_does_not_apply_number_limit_or_increase_early():
537 class _Database:
538 async def get_latest_aweme_time(self, _author_id):
539 return 1700000000
540
541 class _API:
542 async def get_user_collect_mix(self, _sec_uid, max_cursor=0, count=20):
543 return {
544 "items": [
545 {"mix_info": {"mix_id": "mix-1"}},
546 {"mix_info": {"mix_id": "mix-2"}},
547 ],
548 "has_more": False,
549 "max_cursor": 0,
550 }
551
552 async def get_mix_aweme(self, mix_id, cursor=0, count=20):
553 if mix_id == "mix-1":
554 return {
555 "items": [{"aweme_id": "mix-aweme-1", "create_time": 1700000001}],
556 "has_more": False,
557 "max_cursor": 0,
558 }
559 return {
560 "items": [{"aweme_id": "mix-aweme-2", "create_time": 1700000002}],
561 "has_more": False,
562 "max_cursor": 0,
563 }
564
565 class _Downloader:
566 def __init__(self):
567 self.api_client = _API()
568 self.rate_limiter = _NoopRateLimiter()
569 self.database = _Database()
570 self.config = type(
571 "Cfg",
572 (),
573 {
574 "get": lambda _self, key, default=None: {
575 "number": {"collectmix": 1},
576 "increase": {"collectmix": True},
577 }.get(key, default)
578 },
579 )()
580 self._filter_by_time = lambda items: items
581 self._limit_count = lambda items, _mode: items
582
583 strategy = CollectMixUserModeStrategy(_Downloader())
584 items = asyncio.run(strategy.collect_items("self", {"uid": "self"}))
585 assert [item["aweme_id"] for item in items] == ["mix-aweme-1", "mix-aweme-2"]
586
586 lines PYTHON