from __future__ import annotations import httpx import acquisition.classify as classify_module import acquisition.creation_search as creation_search_module from acquisition import store from acquisition.classify import _providers, classify_imgtext, classify_video from acquisition.creation_search import Candidate, PlatformRateLimiter, run_platform_query from acquisition.crawler import RateLimiter from core.models import Post from core.config import PgConfig, Settings from scripts import run_creation_search as run_cli def _settings() -> Settings: return Settings( pg=PgConfig(host="h", port=5432, user="u", password="p", database="d"), aiddit_crawler_base_url="http://crawler.test", crawler_timeout=30, openrouter_timeout_seconds=90, openrouter_model="m", openrouter_base_url="http://openrouter.test", openrouter_api_key="k", llm_model="m", max_cards=12, frames_dir="f", douyin_ratio="540p", data_dir="data", ) def test_creation_store_summary_and_detail(tmp_path): c = store.connect(tmp_path / "app.db") store.create_creation_run(c, "r1", total_queries=1, ts=1) store.ensure_creation_job(c, "r1", "q", "douyin", ts=1) item_id = store.upsert_creation_item( c, run_id="r1", query="q", platform="douyin", rank=1, source_id="dy1", url="https://www.douyin.com/video/1", title="标题", cover_url="https://cover.test/a.jpg", video_url="https://cdn.test/v.mp4", is_displayable=True, ts=2, ) store.upsert_creation_classification( c, item_id, 1, reason="教你做视频", knowledge="先选题再写脚本", prompt_version="abc", ts=3 ) xhs_creation = store.upsert_creation_item( c, run_id="r1", query="q", platform="xiaohongshu", rank=1, title="小红书创作帖", image_urls=["/data/x1.jpg"], is_displayable=True, ts=3, ) xhs_other = store.upsert_creation_item( c, run_id="r1", query="q", platform="xiaohongshu", rank=2, title="小红书非创作帖", image_urls=["/data/x2.jpg"], is_displayable=True, ts=3, ) xhs_hidden = store.upsert_creation_item( c, run_id="r1", query="q", platform="xiaohongshu", rank=3, title="非展示候选", image_urls=["/data/x3.jpg"], is_displayable=False, ts=3, ) wx_creation = store.upsert_creation_item( c, run_id="r1", query="q", platform="weixin", rank=1, title="公众号创作帖", image_urls=["/data/w1.jpg"], is_displayable=True, ts=3, ) store.upsert_creation_classification(c, xhs_creation, 1, reason="ok", knowledge="k", ts=3) store.upsert_creation_classification(c, xhs_other, 0, reason="no", knowledge="", ts=3) store.upsert_creation_classification(c, xhs_hidden, 1, reason="hidden", knowledge="hidden", ts=3) store.upsert_creation_classification(c, wx_creation, 1, reason="ok", knowledge="k", ts=3) store.update_creation_job( c, "r1", "q", "douyin", status="done", attempts=1, searched_count=10, display_count=1, ts=4 ) summary = store.creation_search_summary(c) assert summary["run_id"] == "r1" q_summary = summary["queries"]["q"] assert q_summary["platforms"]["douyin"]["display_count"] == 1 assert q_summary["imgtext_creation_count"] == 2 assert q_summary["imgtext_classified_count"] == 3 assert q_summary["imgtext_total_count"] == 3 assert q_summary["imgtext_target_count"] == 10 detail = store.get_creation_query_detail(c, "q") item = detail["platforms"]["douyin"]["items"][0] assert item["video_url"] == "https://cdn.test/v.mp4" assert item["classification"]["is_creation"] == 1 assert item["classification"]["knowledge"] == "先选题再写脚本" def test_platform_rate_limiter_shares_keyword_and_detail_bucket(): now = {"t": 0.0} sleeps = [] def now_fn(): return now["t"] def sleep_fn(seconds): sleeps.append(seconds) now["t"] += seconds base = RateLimiter( min_interval_seconds=10.0, max_interval_seconds=10.0, now_fn=now_fn, sleep_fn=sleep_fn, ) gate = PlatformRateLimiter("douyin", delegate=base) gate.wait("douyin_keyword") now["t"] += 1.0 gate.wait("douyin_detail") assert sleeps == [9.0] def test_classify_video_accepts_cdn_url(monkeypatch): seen = {} def fake_judge(messages, settings, timeout): seen["url"] = messages[1]["content"][1]["video_url"]["url"] return 1, "ok", "knowledge", "" monkeypatch.setattr("acquisition.classify._judge", fake_judge) res = classify_video({"video": "https://cdn.test/video.mp4"}, _settings()) assert res == (1, "ok", "knowledge", "") assert seen["url"] == "https://cdn.test/video.mp4" def test_classify_imgtext_maps_public_data_to_settings_data_dir(tmp_path, monkeypatch): legacy = tmp_path / "legacy_data" img = legacy / "media" / "x.jpg" img.parent.mkdir(parents=True) img.write_bytes(b"\xff\xd8fake-jpeg") settings = _settings() settings.data_dir = str(legacy) seen = {} def fake_judge(messages, settings, timeout): seen["parts"] = messages[1]["content"] return 1, "ok", "knowledge", "" monkeypatch.setattr("acquisition.classify._judge", fake_judge) res = classify_imgtext({"images": ["/data/media/x.jpg"]}, settings) assert res == (1, "ok", "knowledge", "") image_parts = [p for p in seen["parts"] if p["type"] == "image_url"] assert image_parts and image_parts[0]["image_url"]["url"].startswith("data:image/jpeg;base64,") def test_qwen_provider_uses_dashscope_credentials(monkeypatch): monkeypatch.setenv("CLASSIFY_PROVIDER", "qwen") monkeypatch.setenv("CLASSIFY_MODEL", "qwen3.7-plus") monkeypatch.setenv("ALIYUN_BAILIAN_API_KEY", "sk-test") monkeypatch.setenv("ALIYUN_BAILIAN_BASE_URL", "https://dashscope.test/compatible-mode/v1") monkeypatch.setenv("OPENROUTER_MODEL", "google/gemini-3-flash-preview") providers = _providers(_settings(), [{"role": "user", "content": "ping"}]) assert [p[0] for p in providers] == ["qwen:qwen3.7-plus", "qwen:qwen-vl-plus"] assert providers[0][1] == "https://dashscope.test/compatible-mode/v1/chat/completions" assert providers[0][3]["model"] == "qwen3.7-plus" def test_provider_limiter_and_429_backoff_are_provider_scoped(monkeypatch): assert classify_module._provider_min_interval("qwen", {"CLASSIFY_QWEN_MIN_INTERVAL_SECONDS": "0.25"}) == 0.25 assert classify_module._provider_min_interval("ark", {"CLASSIFY_ARK_MIN_INTERVAL_SECONDS": "0.75"}) == 0.75 now = {"t": 0.0} sleeps = [] calls = {"n": 0} classify_module._PROVIDER_THROTTLES.clear() monkeypatch.setenv("CLASSIFY_PROVIDER", "qwen") monkeypatch.setenv("CLASSIFY_MODEL", "qwen3.7-plus") monkeypatch.setenv("ALIYUN_BAILIAN_API_KEY", "sk-test") monkeypatch.setenv("ALIYUN_BAILIAN_BASE_URL", "https://dashscope.test/compatible-mode/v1") monkeypatch.setenv("CLASSIFY_QWEN_MIN_INTERVAL_SECONDS", "0") monkeypatch.setenv("CLASSIFY_QWEN_429_BACKOFF_SECONDS", "9,30,90") monkeypatch.setattr(classify_module.time, "monotonic", lambda: now["t"]) def fake_sleep(seconds): sleeps.append(seconds) now["t"] += seconds def fake_post(*args, **kwargs): calls["n"] += 1 if calls["n"] == 1: return httpx.Response(429, request=httpx.Request("POST", "https://dashscope.test")) return httpx.Response( 200, json={"choices": [{"message": {"content": '{"is_empty": false, "reason": "ok", "knowledge": "k"}'}}]}, request=httpx.Request("POST", "https://dashscope.test"), ) monkeypatch.setattr(classify_module.time, "sleep", fake_sleep) monkeypatch.setattr(classify_module.httpx, "post", fake_post) result = classify_module._judge([{"role": "user", "content": "ping"}], _settings(), timeout=1) assert result == (1, "ok", "k", "") assert calls["n"] == 2 assert sleeps == [2, 7.0] def test_run_platform_query_skips_bad_items_and_fills_display_limit(tmp_path, monkeypatch): c = store.connect(tmp_path / "app.db") store.create_creation_run(c, "r1", total_queries=1, ts=1) store.ensure_creation_job(c, "r1", "q", "weixin", ts=1) def fake_search(platform, query, *, settings, limit, rate_limiter): assert platform == "weixin" return [Candidate(rank=i, url=f"https://mp.test/{i}", title=f"t{i}") for i in range(1, 8)] def fake_process(platform, candidate, *, query_hash, settings, rate_limiter, downloader): if candidate.rank == 1: raise RuntimeError("detail failed") return { "source_id": str(candidate.rank), "url": candidate.url, "title": candidate.title, "body_text": "正文", "cover_url": "/data/a.jpg", "image_urls": ["/data/a.jpg"], "video_url": "", "raw": {}, } monkeypatch.setattr("acquisition.creation_search.search_candidates", fake_search) monkeypatch.setattr("acquisition.creation_search.process_candidate", fake_process) monkeypatch.setattr("acquisition.creation_search._classify_and_store", lambda *a, **k: None) res = run_platform_query( c, run_id="r1", query="q", platform="weixin", settings=_settings(), search_limit=7, display_limit=5, classify=True, sleep_fn=lambda _: None, ) assert res["status"] == "done" assert res["display_count"] == 5 detail = store.get_creation_query_detail(c, "q") assert len(detail["platforms"]["weixin"]["items"]) == 5 def test_xhs_media_saved_under_settings_data_dir(tmp_path, monkeypatch): settings = _settings() settings.data_dir = str(tmp_path / "legacy_data") post = Post( id="xhs_1", platform="xiaohongshu", url="https://xhs.test/1", content_id="cid1", title="t", body_text="b", image_urls=["https://img.test/a.jpg"], ) monkeypatch.setattr(creation_search_module, "fetch_post_detail", lambda *a, **k: post) out = creation_search_module._process_xhs( Candidate(rank=1, source_id="cid1"), query_hash="qh", settings=settings, rate_limiter=PlatformRateLimiter("xiaohongshu", delegate=RateLimiter( min_interval_seconds=0, max_interval_seconds=0, sleep_fn=lambda _: None, )), downloader=lambda url, platform: b"\xff\xd8fake-jpeg", ) expected = tmp_path / "legacy_data" / "media" / "xiaohongshu" / "qh" / "cid1" / "image_1.jpg" assert expected.exists() assert out["image_urls"] == ["/data/media/xiaohongshu/qh/cid1/image_1.jpg"] def test_creation_job_done_and_classification_queue(tmp_path): c = store.connect(tmp_path / "app.db") store.create_creation_run(c, "r1", total_queries=1, ts=1) store.ensure_creation_job(c, "r1", "q", "weixin", ts=1) assert not store.creation_job_is_done(c, "r1", "q", "weixin", display_limit=5) store.update_creation_job(c, "r1", "q", "weixin", status="done", display_count=5, ts=2) assert store.creation_job_is_done(c, "r1", "q", "weixin", display_limit=5) missing = store.upsert_creation_item( c, run_id="r1", query="q", platform="weixin", rank=1, title="missing", image_urls=["/data/a.jpg"], is_displayable=True, ts=3, ) failed = store.upsert_creation_item( c, run_id="r1", query="q", platform="weixin", rank=2, title="failed", image_urls=["/data/b.jpg"], is_displayable=True, ts=3, ) done = store.upsert_creation_item( c, run_id="r1", query="q", platform="weixin", rank=3, title="done", image_urls=["/data/c.jpg"], is_displayable=True, ts=3, ) store.upsert_creation_classification(c, failed, None, reason="bad", error="bad", ts=4) store.upsert_creation_classification(c, done, 1, reason="ok", knowledge="k", ts=4) retry_rows = store.creation_items_to_classify(c, run_id="r1", platforms=["weixin"], retry_failed=True) missing_rows = store.creation_items_to_classify(c, run_id="r1", platforms=["weixin"], retry_failed=False) assert {r["id"] for r in retry_rows} == {missing, failed} assert {r["id"] for r in missing_rows} == {missing} def test_run_workers_isolates_worker_exception(tmp_path, monkeypatch): db = tmp_path / "app.db" c = store.connect(db) store.create_creation_run(c, "r1", total_queries=1, ts=1) for p in ["xiaohongshu", "weixin"]: store.ensure_creation_job(c, "r1", "q", p, ts=1) c.close() def fake_run(conn, *, platform, run_id, query, display_limit, **kwargs): if platform == "xiaohongshu": raise RuntimeError("boom") store.update_creation_job( conn, run_id, query, platform, status="done", searched_count=10, display_count=display_limit, ts=2, ) return {"platform": platform, "status": "done", "display_count": display_limit, "error": ""} monkeypatch.setattr(run_cli, "run_platform_query", fake_run) failed = run_cli.run_platform_workers( run_id="r1", queries=["q"], platforms=["xiaohongshu", "weixin"], settings=_settings(), search_limit=10, display_limit=5, classify=False, skip_done=False, db_path=db, ) c = store.connect(db) summary = store.creation_search_summary(c, run_id="r1")["queries"]["q"]["platforms"] assert failed == 1 assert summary["xiaohongshu"]["status"] == "failed" assert "worker异常" in summary["xiaohongshu"]["error"] assert summary["weixin"]["status"] == "done" def test_run_workers_skip_done_does_not_rerun(tmp_path, monkeypatch): db = tmp_path / "app.db" c = store.connect(db) store.create_creation_run(c, "r1", total_queries=1, ts=1) store.ensure_creation_job(c, "r1", "q", "weixin", ts=1) store.update_creation_job(c, "r1", "q", "weixin", status="done", display_count=5, ts=2) c.close() calls = [] def fake_run(*args, **kwargs): calls.append(kwargs) return {"platform": "weixin", "status": "done", "display_count": 5, "error": ""} monkeypatch.setattr(run_cli, "run_platform_query", fake_run) failed = run_cli.run_platform_workers( run_id="r1", queries=["q"], platforms=["weixin"], settings=_settings(), search_limit=10, display_limit=5, classify=False, skip_done=True, db_path=db, ) assert failed == 0 assert calls == [] def test_run_workers_resolves_env_db_path(tmp_path, monkeypatch): db = tmp_path / "legacy_data" / "app.db" c = store.connect(db) store.create_creation_run(c, "r1", total_queries=1, ts=1) store.ensure_creation_job(c, "r1", "q", "weixin", ts=1) c.close() monkeypatch.setenv("CK_SQLITE_PATH", str(db)) def fake_run(conn, *, platform, run_id, query, display_limit, **kwargs): store.update_creation_job( conn, run_id, query, platform, status="done", searched_count=10, display_count=display_limit, ts=2, ) return {"platform": platform, "status": "done", "display_count": display_limit, "error": ""} monkeypatch.setattr(run_cli, "run_platform_query", fake_run) failed = run_cli.run_platform_workers( run_id="r1", queries=["q"], platforms=["weixin"], settings=_settings(), search_limit=10, display_limit=5, classify=False, skip_done=False, ) c = store.connect(db) summary = store.creation_search_summary(c, run_id="r1")["queries"]["q"]["platforms"] assert failed == 0 assert summary["weixin"]["status"] == "done"