| 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 |