Create-knowledge-find-decode/tests/test_creation_search.py

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"