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

"""M3 离线测试：注入假 http_post，验证消息构造 + JSON 解析（不调真实 Gemini）。

真实 Gemini 验收见 scripts/smoke_extract.py（须在能联网调 OpenRouter 的云端跑）。
"""
from __future__ import annotations

import json

from creation_knowledge.integrations.extractor import GeminiExtractor
from core.models import Post


class _FakeResp:
    def __init__(self, payload: dict):
        self._payload = payload

    def raise_for_status(self):
        return None

    def json(self):
        return self._payload


def _fake_post_returning(content: str):
    captured = {}

    def _post(url, headers=None, json=None, timeout=None):
        captured["url"] = url
        captured["headers"] = headers
        captured["json"] = json
        return _FakeResp({"choices": [{"message": {"content": content}}]})

    return _post, captured


def _post(**kw) -> Post:
    return Post(
        id="xhs_x", url="u", content_id="x",
        title="只要学会这几样，写短视频脚本真的不难",
        body_text="", topic_list=["短视频"],
        image_urls=["https://img/1.webp", "https://img/2.webp"],
        **kw,
    )


def test_build_messages_has_text_and_images():
    client = GeminiExtractor(api_key="k", http_post=lambda **kw: None)
    msgs = client.build_messages(_post())
    assert msgs[0]["role"] == "system"
    user_parts = msgs[1]["content"]
    assert user_parts[0]["type"] == "text"
    images = [p for p in user_parts if p["type"] == "image_url"]
    assert len(images) == 2
    assert images[0]["image_url"]["url"] == "https://img/1.webp"


def test_parse_plain_json():
    fake, captured = _fake_post_returning(
        json.dumps({"text": "脚本要素：标题/地点/分镜", "from_image": "九宫格讲了6要素",
                    "from_video": "", "is_empty": False}, ensure_ascii=False)
    )
    client = GeminiExtractor(api_key="k", http_post=fake)
    out = client.extract(_post())
    assert out.is_empty is False
    assert "from_image" and out.from_image
    assert "脚本" in out.text
    # 鉴权头 + 模型名正确传入
    assert captured["headers"]["Authorization"].startswith("Bearer ")
    assert captured["json"]["model"]


def test_build_messages_labels_each_card():
    client = GeminiExtractor(api_key="k", http_post=lambda **kw: None)
    parts = client.build_messages(_post())[1]["content"]
    labels = [p["text"] for p in parts if p["type"] == "text"]
    assert any("【卡片1】" in t for t in labels)
    assert any("【卡片2】" in t for t in labels)


def test_parse_per_card_output():
    fake, _ = _fake_post_returning(json.dumps(
        {"text": "x", "cards": [{"index": 1, "content": "卡1要点"},
                                {"index": 2, "content": "卡2要点"}],
         "is_empty": False}, ensure_ascii=False))
    out = GeminiExtractor(api_key="k", http_post=fake).extract(_post())
    assert len(out.cards) == 2
    assert out.cards[0].index == 1 and out.cards[0].content == "卡1要点"


def test_parse_json_with_code_fence():
    fenced = "```json\n" + json.dumps(
        {"text": "x", "from_image": "y", "from_video": "", "is_empty": "false"},
        ensure_ascii=False) + "\n```"
    fake, _ = _fake_post_returning(fenced)
    client = GeminiExtractor(api_key="k", http_post=fake)
    out = client.extract(_post())
    assert out.is_empty is False
    assert out.from_image == "y"