返回 douyin-downloader
test_discovery.py
根目录 / tests / test_discovery.py
1 """热榜 / 搜索落盘模块测试。"""
2
3 import json
4 from pathlib import Path
5 from typing import Any, Dict, List
6
7 import pytest
8
9 from core.discovery import dump_hot_board, search_and_dump
10
11
12 class _FakeAPIClient:
13 def __init__(
14 self,
15 hot_items: List[Dict[str, Any]] | None = None,
16 search_pages: List[Dict[str, Any]] | None = None,
17 ):
18 self._hot_items = hot_items or []
19 self._search_pages = list(search_pages or [])
20 self.search_calls: List[Dict[str, Any]] = []
21
22 async def get_hot_search_board(self) -> Dict[str, Any]:
23 return {
24 "items": self._hot_items,
25 "has_more": False,
26 "max_cursor": 0,
27 }
28
29 async def search_aweme(self, keyword, *, offset, count, sort_type=0, publish_time=0):
30 self.search_calls.append({"keyword": keyword, "offset": offset, "count": count})
31 if not self._search_pages:
32 return {"items": [], "has_more": False, "max_cursor": offset}
33 return self._search_pages.pop(0)
34
35
36 @pytest.mark.asyncio
37 async def test_dump_hot_board_writes_jsonl(tmp_path):
38 api = _FakeAPIClient(
39 hot_items=[{"word": "foo", "hot_value": 100}, {"word": "bar", "hot_value": 50}]
40 )
41 result = await dump_hot_board(api, tmp_path)
42 assert result["count"] == 2
43 out = Path(result["path"])
44 assert out.exists()
45 lines = out.read_text(encoding="utf-8").strip().splitlines()
46 assert len(lines) == 2
47 assert json.loads(lines[0])["word"] == "foo"
48
49
50 @pytest.mark.asyncio
51 async def test_dump_hot_board_respects_limit(tmp_path):
52 api = _FakeAPIClient(hot_items=[{"word": f"w{i}"} for i in range(20)])
53 result = await dump_hot_board(api, tmp_path, limit=5)
54 assert result["count"] == 5
55
56
57 @pytest.mark.asyncio
58 async def test_search_and_dump_accumulates_pages(tmp_path):
59 api = _FakeAPIClient(
60 search_pages=[
61 {
62 "items": [{"aweme_id": "1"}, {"aweme_id": "2"}],
63 "has_more": True,
64 "max_cursor": 2,
65 },
66 {
67 "items": [{"aweme_id": "3"}, {"aweme_id": "2"}], # dup
68 "has_more": False,
69 "max_cursor": 4,
70 },
71 ]
72 )
73 result = await search_and_dump(api, "cat", tmp_path, max_items=0)
74 assert result["count"] == 3
75 assert {"1", "2", "3"} == {c["aweme_id"] for c in result["items"]}
76
77
78 @pytest.mark.asyncio
79 async def test_search_respects_max_items(tmp_path):
80 api = _FakeAPIClient(
81 search_pages=[
82 {"items": [{"aweme_id": str(i)} for i in range(5)], "has_more": True, "max_cursor": 5},
83 {
84 "items": [{"aweme_id": str(i)} for i in range(5, 10)],
85 "has_more": False,
86 "max_cursor": 10,
87 },
88 ]
89 )
90 result = await search_and_dump(api, "cat", tmp_path, max_items=3, page_size=5)
91 assert result["count"] == 3
92
93
94 @pytest.mark.asyncio
95 async def test_search_stops_on_stuck_cursor(tmp_path):
96 api = _FakeAPIClient(
97 search_pages=[
98 {"items": [{"aweme_id": "1"}], "has_more": True, "max_cursor": 0},
99 ]
100 * 10
101 )
102 result = await search_and_dump(api, "cat", tmp_path, max_items=0)
103 # 第一页 cursor 未推进应立即停止
104 assert len(api.search_calls) == 1
105 assert result["count"] == 1
106
106 lines PYTHON